v0.3.0 post-implementation audit: resolve all 9 tracked discrepancies

Works through docs/v0.3.0-followups.md item by item, closing the gap
between the design doc and the shipped v0.3.0-stage2-autoregressive code:

1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2
   dispatch, stage-2 particle-type-class marginal.
2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run.
3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/
   train instead of the hardcoded K_MAX constant.
4. Mixed conditioning.particle.type / conditioning.material.type support
   end-to-end (data pipeline + dwarf warm-cache).
5. conditioning.share_stages = true: one shared ConditionEncoder instance
   across both stages.
6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2
   (was silently unimplemented).
7. giant predict/rollout: implement conditioning.*.type = "onehot" via the
   checkpoint's saved pdg_topn_map/mat_topn_map.
8. network.py's checkpoint-path model_config migration now fails loudly on
   non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's
   TOML-load path (§4.2).
9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a
   rollout-capable checkpoint (§9).

Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics),
mostly a test-helper dict-unpack pattern that made every unrelated
constructor keyword look like a type error.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

docs/v0.3.0-design.md

@@ -892,6 +892,13 @@ All in `giant/config.py`:
   undesigned:** the mechanism should be worked out on its own terms when it is
   taken up, not pre-shaped by choices made for this refactor. Adding the block
   later is a config addition, not a break.
+- **`stage2_model.generator = "ddpm"`.** The value is **accepted by the
+  schema** (§3.3 lists `"flow" | "ddpm" | "wgan"` with no caveat) but
+  `FlowDDPMStageTrainer.__init__` (`giant/train.py`) raises
+  `NotImplementedError` for stage 2 — only `"flow"` and `"wgan"` have a
+  stage-2 secondary-decoder loss implemented. `stage1_model.generator =
+  "ddpm"` is unaffected; this restriction is stage-2-only. Landing stage-2
+  ddpm later is a trainer addition, not a config break.
 
 ### 11.3 Tracked as implementation work
 

docs/v0.3.0-followups.md

@@ -0,0 +1,37 @@
+# v0.3.0 — post-implementation audit: open discrepancies
+
+**Status:** steps 1–7 of `docs/v0.3.0-design.md` §12 are implemented (branch
+`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document
+tracked discrepancies found between that implementation and the design contract
+during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now
+resolved** — either implemented (1-5, 7-9) or explicitly deferred into
+`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8
+(`estimate_batch_size` recalibration) was intentionally still outstanding per
+§12 and was never tracked here.
+
+---
+
+## Confirmed correct during the audit (no action needed)
+
+For reference — these were explicitly checked against the design doc and
+match it, including two spots the doc itself flagged as likely stale that
+turned out fine:
+
+- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/
+  `merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum,
+  `n_sec.mode = "stop_token"` error (§9, §11.2).
+- `network.py`'s full class decomposition (§5.3), dict-returning
+  `build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims,
+  ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history
+  encoders (§6.2).
+- Shared PDG top-N type map (one map, not two, per §8), `other_policy`
+  sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance
+  diagnostic surfaced in `giant analyze` (§11.3).
+- `analysis/render.py`/`analysis/router_gating.py` correctly branch
+  old-flat vs new-nested `model_config["router"]` location — doc flagged this
+  as a likely stale spot (§10) but it's actually fine.
+- `train.py`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
+  per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only
+  training via ground-truth `x1_s1` (§7).
+- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init`
+  seeding, removed stale expert-size warning (§9, §10).

giant/analysis/router_gating.py

@@ -61,21 +61,26 @@ class _RouterHandle:
     pdg_map: dict[int, int]
     mat_map: dict[str, int]
     cond_normalizer: "Normalizer"
-    conditioning: str
+    particle_conditioning: str
+    material_conditioning: str
     router_type: str
 
 
-def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str:
-    """The single conditioning-mode string `giant.data.transforms.
-    build_cond_features` still expects — from either a v0.2 checkpoint's
-    flat `model_config["conditioning"]` (already a string) or a new-format
-    one (`model_config["conditioning"]["particle"]["type"]`; `validate_config`
-    guarantees particle/material agree, see its mixed-conditioning check).
-    Mirrors `giant.cli._conditioning_str`."""
+def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features` — from either a v0.2
+    checkpoint's flat `model_config["conditioning"]` (one shared string, same
+    for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` —
+    docs/v0.3.0-design.md §3.1: the two axes may differ). Mirrors
+    `giant.cli._conditioning_axes`."""
     raw = model_cfg.get("conditioning", default)
     if isinstance(raw, dict):
-        return raw.get("particle", {}).get("type", default)
-    return raw
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
+        )
+    return raw, raw
 
 
 def load_router(checkpoint: str | Path) -> _RouterHandle | None:
@@ -107,12 +112,14 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     if router is None:
         return None
 
+    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
     return _RouterHandle(
         router=router,
         pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},
         mat_map={str(k): v for k, v in ckpt["mat_map"].items()},
         cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]),
-        conditioning=_conditioning_str(model_cfg),
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         router_type=router_cfg["type"],
     )
 
@@ -169,7 +176,8 @@ def _gate_for_df(
         handle.pdg_map,
         handle.mat_map,
         cond_normalizer=handle.cond_normalizer,
-        conditioning=handle.conditioning,
+        particle_conditioning=handle.particle_conditioning,
+        material_conditioning=handle.material_conditioning,
     )
     with torch.no_grad():
         gate = handle.router.gate(

giant/cli.py

@@ -19,6 +19,7 @@ from tqdm import tqdm
 
 from giant import config as gconfig
 from giant.constants import (
+    K_MAX,
     LOCAL_TARGET_NAMES,
     PREDICT_COORD_METADATA_KEY,
     PREDICT_SCHEMA_VERSION,
@@ -72,44 +73,21 @@ def _router_total_experts(router_cfg: dict) -> int:
     return int(router_cfg.get("n_experts", 1))
 
 
-def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str:
-    """The single conditioning-mode string `giant.data.transforms.
-    build_cond_features` still expects — from either a v0.2 checkpoint's
-    flat `model_config["conditioning"]` (already a string) or a new-format
-    one (`model_config["conditioning"]["particle"]["type"]`; `validate_config`
-    guarantees particle/material agree, see its mixed-conditioning check)."""
+def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features`/`build_features` — from
+    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
+    shared string, same for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` —
+    docs/v0.3.0-design.md §3.1: the two axes are configured independently and
+    may differ)."""
     raw = model_cfg.get("conditioning", default)
     if isinstance(raw, dict):
-        return raw.get("particle", {}).get("type", default)
-    return raw
-
-
-def _check_conditioning_onehot_support(model_cfg: dict, command: str) -> None:
-    """`giant predict`/`giant rollout` don't yet support
-    `conditioning.particle/material.type = "onehot"` — that needs
-    `pdg_topn_map.class_map`/`mat_topn_map.class_map` threaded into every
-    `build_cond_features` call in this module and `giant/rollout.py`'s
-    `_step_chunk`, which hasn't been wired (a separate axis from
-    `stage2_model.particle_type.target`, handled below — see
-    docs/v0.3.0-design.md §3.1 vs §3.3). Without this guard, predict/rollout
-    would silently build a `cond_cat` missing the topN columns
-    `ConditionEncoder`'s `"onehot"` mode expects. A v0.2 (flat) `model_config`
-    never has this, so this is a no-op there.
-    """
-    conditioning = model_cfg.get("conditioning")
-    if not isinstance(conditioning, dict):
-        return
-    particle_type = conditioning.get("particle", {}).get("type")
-    material_type = conditioning.get("material", {}).get("type")
-    if particle_type == "onehot" or material_type == "onehot":
-        typer.echo(
-            f"error: giant {command} does not yet support "
-            "conditioning.particle/material.type = 'onehot' — only "
-            "stage2_model.particle_type.target in ('onehot', 'embedding') "
-            "is implemented (docs/v0.3.0-design.md step 6).",
-            err=True,
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
         )
-        raise typer.Exit(1)
+    return raw, raw
 
 
 def _stage_cfg(model_cfg: dict, stage: str) -> dict:
@@ -144,6 +122,14 @@ def _load_pdg_topn_map(ckpt: dict):
     return topnmap_from_json(raw, axis="pdg") if raw is not None else None
 
 
+def _load_mat_topn_map(ckpt: dict):
+    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
+    `giant.pipeline.run_setup_stage`)."""
+    raw = ckpt.get("mat_topn_map")
+    return topnmap_from_json(raw, axis="material") if raw is not None else None
+
+
 def _batch_size_estimate_dims(
     model_cfg: dict, training: bool, stage: str = "stage1"
 ) -> tuple[int, int]:
@@ -295,6 +281,16 @@ class Mode(str, Enum):
     wgan = "wgan"
 
 
+class Decoder(str, Enum):
+    one_shot = "one_shot"
+    autoregressive = "autoregressive"
+
+
+class Stage1Context(str, Enum):
+    truth = "truth"
+    sampled = "sampled"
+
+
 # Conditioning itself lives in giant.config (imported below as gconfig) —
 # shared with scripts/dwarf.py's Typer commands so the two CLIs can't
 # silently drift apart on the option's valid values.
@@ -389,6 +385,89 @@ def train(
             "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
         ),
     ] = None,
+    stage1_generator: Annotated[
+        Optional[Mode],
+        typer.Option(
+            "--stage1-generator",
+            help="Stage 1's generative objective — overrides --mode for stage 1 "
+            "only (see docs/v0.3.0-design.md decision 7)",
+        ),
+    ] = None,
+    stage1_hidden_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-hidden-dim",
+            help="Overrides --hidden-dim for stage 1 only (same effect today; "
+            "--hidden-dim is kept as a shorthand since stage 1 was the only "
+            "target before stage2_model got its own flags)",
+        ),
+    ] = None,
+    stage1_n_res_blocks: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"
+        ),
+    ] = None,
+    stage1_dropout: Annotated[
+        Optional[float],
+        typer.Option("--stage1-dropout", help="Overrides --dropout for stage 1 only"),
+    ] = None,
+    stage2_generator: Annotated[
+        Optional[Mode],
+        typer.Option(
+            "--stage2-generator",
+            help="Stage 2's generative objective — overrides --mode for stage 2 "
+            "only, e.g. combine with --stage1-generator flow for a mixed "
+            "flow/wgan run",
+        ),
+    ] = None,
+    stage2_hidden_dim: Annotated[
+        Optional[int],
+        typer.Option("--stage2-hidden-dim", help="Stage 2 trunk width"),
+    ] = None,
+    stage2_n_res_blocks: Annotated[
+        Optional[int],
+        typer.Option("--stage2-n-res-blocks", help="Stage 2 trunk depth"),
+    ] = None,
+    stage2_dropout: Annotated[
+        Optional[float],
+        typer.Option("--stage2-dropout", help="Dropout inside stage 2's ResBlocks"),
+    ] = None,
+    stage2_decoder: Annotated[
+        Optional[Decoder],
+        typer.Option(
+            "--stage2-decoder",
+            help="one_shot: predict all k_max secondary slots at once (v0.2 "
+            "behaviour). autoregressive: emit one secondary at a time in "
+            "descending-energy order (default)",
+        ),
+    ] = None,
+    stage2_k_max: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-k-max",
+            help="Maximum secondary slots (fixed width under one_shot, a "
+            "generation-loop safety cap under autoregressive; default: 15)",
+        ),
+    ] = None,
+    stage2_context_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-context-dim",
+            help="Width of the projected stage-1 outcome fed into stage 2's "
+            "conditioning (default: 64)",
+        ),
+    ] = None,
+    stage2_stage1_context: Annotated[
+        Optional[Stage1Context],
+        typer.Option(
+            "--stage2-stage1-context",
+            help="What stage 2 conditions on during training: 'truth' (the "
+            "ground-truth stage-1 target, detached — default) or 'sampled' "
+            "(stage 1's own sampled output, closing the train/inference gap "
+            "at the cost of an extra sampling pass per batch)",
+        ),
+    ] = None,
     conditioning: Annotated[
         Optional[Conditioning],
         typer.Option(
@@ -457,6 +536,52 @@ def train(
             "(default: same as --lr)",
         ),
     ] = None,
+    stage1_n_critic: Annotated[
+        Optional[int],
+        typer.Option("--stage1-n-critic", help="Overrides --n-critic for stage 1 only"),
+    ] = None,
+    stage1_gp_weight: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"
+        ),
+    ] = None,
+    stage1_noise_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"
+        ),
+    ] = None,
+    stage1_critic_lr: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"
+        ),
+    ] = None,
+    stage2_n_critic: Annotated[
+        Optional[int],
+        typer.Option("--stage2-n-critic", help="Overrides --n-critic for stage 2 only"),
+    ] = None,
+    stage2_gp_weight: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"
+        ),
+    ] = None,
+    stage2_noise_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-noise-dim",
+            help="Overrides --noise-dim for stage 2 only; under "
+            "--stage2-decoder autoregressive a fresh draw is made per token",
+        ),
+    ] = None,
+    stage2_critic_lr: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"
+        ),
+    ] = None,
     val_fraction: Annotated[
         Optional[float], typer.Option("--val-fraction", "-f")
     ] = None,
@@ -597,9 +722,11 @@ def train(
         }.items()
         if v is not None
     }
-    # Architecture-shape flags apply to stage1_model only (decision: stage 2
-    # is tuned via a hand-edited config.toml until stage-prefixed --stage2-*
-    # flags land — see docs/v0.3.0-design.md decision 7).
+    # --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for
+    # backward compatibility (they predate stage2_model having its own
+    # flags); --stage1-*/--stage2-* below are the explicit, discoverable
+    # per-stage flags docs/v0.3.0-design.md decision 7 calls for, and take
+    # precedence when both are given.
     cli_stage1_model: dict[str, object] = {
         k: v
         for k, v in {
@@ -609,6 +736,32 @@ def train(
         }.items()
         if v is not None
     }
+    cli_stage1_model.update(
+        {
+            k: v
+            for k, v in {
+                "hidden_dim": stage1_hidden_dim,
+                "n_res_blocks": stage1_n_res_blocks,
+                "dropout": stage1_dropout,
+            }.items()
+            if v is not None
+        }
+    )
+    cli_stage2_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": stage2_hidden_dim,
+            "n_res_blocks": stage2_n_res_blocks,
+            "dropout": stage2_dropout,
+            "decoder": stage2_decoder.value if stage2_decoder is not None else None,
+            "k_max": stage2_k_max,
+            "context_dim": stage2_context_dim,
+            "stage1_context": stage2_stage1_context.value
+            if stage2_stage1_context is not None
+            else None,
+        }.items()
+        if v is not None
+    }
     cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_stage1_model["router"] = cli_router
@@ -628,37 +781,65 @@ def train(
         overrides["train"] = cli_train
     if cli_stage1_model:
         overrides["stage1_model"] = cli_stage1_model
+    if cli_stage2_model:
+        overrides["stage2_model"] = cli_stage2_model
     if cli_conditioning:
         overrides["conditioning"] = cli_conditioning
 
-    # --mode/--n-critic/--gp-weight/--critic-lr apply to BOTH stages — v0.2
-    # had one global mode/wgan config shared by both (see
-    # giant.config.migrate_config's train.mode /
-    # train.{n_critic,gp_weight,critic_lr} precedent), unlike the
-    # architecture-shape flags above.
+    # --mode/--n-critic/--gp-weight/--critic-lr/--noise-dim apply to BOTH
+    # stages by default (v0.2 had one global mode/wgan config shared by both
+    # — see giant.config.migrate_config's train.mode /
+    # train.{n_critic,gp_weight,critic_lr} precedent); the --stage1-*/
+    # --stage2-* variants below override a single stage independently (decision
+    # 7), which is what actually enables e.g. `--stage1-generator flow
+    # --stage2-generator wgan`.
     if mode is not None:
         overrides.setdefault("stage1_model", {})["generator"] = mode.value
         overrides.setdefault("stage2_model", {})["generator"] = mode.value
-    if noise_dim is not None:
-        overrides.setdefault("stage1_model", {}).setdefault("wgan", {})["noise_dim"] = (
-            noise_dim
-        )
-    wgan_overrides = {
+    if stage1_generator is not None:
+        overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
+    if stage2_generator is not None:
+        overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
+
+    shared_wgan_overrides = {
         k: v
         for k, v in {
             "n_critic": n_critic,
             "gp_weight": gp_weight,
+            "noise_dim": noise_dim,
             "critic_lr": critic_lr,
         }.items()
         if v is not None
     }
-    if wgan_overrides:
-        overrides.setdefault("stage1_model", {}).setdefault("wgan", {}).update(
-            wgan_overrides
-        )
-        overrides.setdefault("stage2_model", {}).setdefault("wgan", {}).update(
-            wgan_overrides
-        )
+    stage1_wgan_overrides = {
+        k: v
+        for k, v in {
+            "n_critic": stage1_n_critic,
+            "gp_weight": stage1_gp_weight,
+            "noise_dim": stage1_noise_dim,
+            "critic_lr": stage1_critic_lr,
+        }.items()
+        if v is not None
+    }
+    stage2_wgan_overrides = {
+        k: v
+        for k, v in {
+            "n_critic": stage2_n_critic,
+            "gp_weight": stage2_gp_weight,
+            "noise_dim": stage2_noise_dim,
+            "critic_lr": stage2_critic_lr,
+        }.items()
+        if v is not None
+    }
+    for stage_name, stage_specific in (
+        ("stage1_model", stage1_wgan_overrides),
+        ("stage2_model", stage2_wgan_overrides),
+    ):
+        stage_wgan = {**shared_wgan_overrides, **stage_specific}
+        if stage_wgan:
+            overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(
+                stage_wgan
+            )
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -735,6 +916,36 @@ def new_run(
     n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
     emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
     dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
+    stage1_generator: Annotated[
+        Optional[Mode], typer.Option("--stage1-generator")
+    ] = None,
+    stage1_hidden_dim: Annotated[
+        Optional[int], typer.Option("--stage1-hidden-dim")
+    ] = None,
+    stage1_n_res_blocks: Annotated[
+        Optional[int], typer.Option("--stage1-n-res-blocks")
+    ] = None,
+    stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None,
+    stage2_generator: Annotated[
+        Optional[Mode], typer.Option("--stage2-generator")
+    ] = None,
+    stage2_hidden_dim: Annotated[
+        Optional[int], typer.Option("--stage2-hidden-dim")
+    ] = None,
+    stage2_n_res_blocks: Annotated[
+        Optional[int], typer.Option("--stage2-n-res-blocks")
+    ] = None,
+    stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None,
+    stage2_decoder: Annotated[
+        Optional[Decoder], typer.Option("--stage2-decoder")
+    ] = None,
+    stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
+    stage2_context_dim: Annotated[
+        Optional[int], typer.Option("--stage2-context-dim")
+    ] = None,
+    stage2_stage1_context: Annotated[
+        Optional[Stage1Context], typer.Option("--stage2-stage1-context")
+    ] = None,
     conditioning: Annotated[
         Optional[Conditioning], typer.Option("--conditioning")
     ] = None,
@@ -802,6 +1013,32 @@ def new_run(
         }.items()
         if v is not None
     }
+    cli_stage1_model.update(
+        {
+            k: v
+            for k, v in {
+                "hidden_dim": stage1_hidden_dim,
+                "n_res_blocks": stage1_n_res_blocks,
+                "dropout": stage1_dropout,
+            }.items()
+            if v is not None
+        }
+    )
+    cli_stage2_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": stage2_hidden_dim,
+            "n_res_blocks": stage2_n_res_blocks,
+            "dropout": stage2_dropout,
+            "decoder": stage2_decoder.value if stage2_decoder is not None else None,
+            "k_max": stage2_k_max,
+            "context_dim": stage2_context_dim,
+            "stage1_context": stage2_stage1_context.value
+            if stage2_stage1_context is not None
+            else None,
+        }.items()
+        if v is not None
+    }
     cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_stage1_model["router"] = cli_router
@@ -818,11 +1055,17 @@ def new_run(
         overrides["train"] = cli_train
     if cli_stage1_model:
         overrides["stage1_model"] = cli_stage1_model
+    if cli_stage2_model:
+        overrides["stage2_model"] = cli_stage2_model
     if cli_conditioning:
         overrides["conditioning"] = cli_conditioning
     if mode is not None:
         overrides.setdefault("stage1_model", {})["generator"] = mode.value
         overrides.setdefault("stage2_model", {})["generator"] = mode.value
+    if stage1_generator is not None:
+        overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
+    if stage2_generator is not None:
+        overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -988,10 +1231,11 @@ def predict(
         raise typer.Exit(1)
 
     model_cfg = ckpt["model_config"]
-    _check_conditioning_onehot_support(model_cfg, "predict")
     pdg_topn_map = _load_pdg_topn_map(ckpt)
+    mat_topn_map = _load_mat_topn_map(ckpt)
     other_policy = _particle_type_other_policy(model_cfg)
     stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
+    stage2_k_max = _stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
 
     if batch_size_auto:
         est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
@@ -1013,7 +1257,21 @@ def predict(
 
     assert batch_size_value is not None
     bs = batch_size_value
-    conditioning = _conditioning_str(model_cfg)
+    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.particle.type='onehot' but has "
+            "no pdg_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.material.type='onehot' but has "
+            "no mat_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
     pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
     mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
     cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
@@ -1050,7 +1308,18 @@ def predict(
     skipped = 0
     unknown_pdg_counts: Counter[int] = Counter()
     total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
-    chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks
+
+    def chunk_iter(path: Path, offset: int):
+        if coord == Coord.local:
+            return iter_file_chunks(path, offset=offset, k_max=stage2_k_max)
+        return iter_cond_chunks(path, offset=offset)
+
+    cond_pdg_topn = (
+        pdg_topn_map.class_map if particle_conditioning == "onehot" else None
+    )
+    cond_mat_topn = (
+        mat_topn_map.class_map if material_conditioning == "onehot" else None
+    )
 
     def _concat(
         a: dict[str, np.ndarray], b: dict[str, np.ndarray]
@@ -1062,12 +1331,26 @@ def predict(
 
         if coord == Coord.local:
             cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
-                piece, pdg_map, mat_map, conditioning=conditioning
+                piece,
+                pdg_map,
+                mat_map,
+                particle_conditioning=particle_conditioning,
+                material_conditioning=material_conditioning,
+                pdg_topn_map=cond_pdg_topn,
+                mat_topn_map=cond_mat_topn,
+                k_max=stage2_k_max,
             )
             cond_cont = cond_norm.transform(cond_cont)
         else:
             cond_cont, cond_cat = build_cond_features(
-                piece, pdg_map, mat_map, cond_norm, conditioning=conditioning
+                piece,
+                pdg_map,
+                mat_map,
+                cond_norm,
+                particle_conditioning=particle_conditioning,
+                material_conditioning=material_conditioning,
+                pdg_topn_map=cond_pdg_topn,
+                mat_topn_map=cond_mat_topn,
             )
 
         cc = torch.from_numpy(cond_cont).float().to(_device)
@@ -1409,9 +1692,23 @@ def rollout(
     training_cfg = gconfig.load_checkpoint_config(checkpoint)
 
     model_cfg = ckpt["model_config"]
-    _check_conditioning_onehot_support(model_cfg, "rollout")
-    conditioning = _conditioning_str(model_cfg)
+    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
     pdg_topn_map = _load_pdg_topn_map(ckpt)
+    mat_topn_map = _load_mat_topn_map(ckpt)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.particle.type='onehot' but has "
+            "no pdg_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.material.type='onehot' but has "
+            "no mat_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
     other_policy = _particle_type_other_policy(model_cfg)
     stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
     stage2_ddpm_steps = _ddpm_steps(model_cfg, "stage2")
@@ -1489,8 +1786,10 @@ def rollout(
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
         on_chunk=_write_chunk,
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,

giant/config.py

@@ -680,18 +680,6 @@ def validate_config(cfg: dict) -> None:
     single block's defaults.
     """
     particle_type = _get_path(cfg, "conditioning.particle.type")
-    material_type = _get_path(cfg, "conditioning.material.type")
-
-    if particle_type != material_type:
-        raise ValueError(
-            "conditioning.particle.type "
-            f"({particle_type!r}) != conditioning.material.type "
-            f"({material_type!r}) — the data pipeline (giant/data/transforms.py's "
-            "build_features/build_cond_features) doesn't support mixed "
-            "particle/material conditioning types yet, even though "
-            "ConditionEncoder itself (giant/model/network.py) is already able "
-            "to. Not implemented in v0.3.0 — use the same type for both axes."
-        )
 
     pt_target = _get_path(cfg, "stage2_model.particle_type.target")
     if pt_target == "embedding" and particle_type != "embedding":
@@ -733,20 +721,34 @@ def validate_config(cfg: dict) -> None:
             "(standalone stage-2 evaluation only, never for rollout)"
         )
 
+    if (
+        _get_path(cfg, "stage2_model.n_sec.mode") == "truth"
+        and _get_path(cfg, "stage1_model.active")
+        and _get_path(cfg, "stage2_model.active")
+    ):
+        raise ValueError(
+            "stage2_model.n_sec.mode = 'truth' is invalid for a "
+            "rollout-capable checkpoint (both stage1_model.active and "
+            "stage2_model.active = true — see docs/v0.3.0-design.md §9): "
+            "'truth' takes n_sec from ground truth, which giant rollout "
+            "doesn't have. 'truth' is for standalone stage-2 evaluation "
+            "only — set stage1_model.active = false for that, or use "
+            "'head' (default) for a rollout-capable checkpoint."
+        )
+
     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'"
             )
 
 

giant/data/dataset.py

@@ -6,6 +6,7 @@ import numpy as np
 import torch
 from torch.utils.data import IterableDataset
 
+from giant.constants import K_MAX
 from giant.data.loader import event_id_offset, iter_file_chunks
 from giant.data.transforms import Normalizer, build_features, sorted_membership
 
@@ -46,7 +47,7 @@ class StreamingStepsDataset(IterableDataset):
                     (see docs/v0.3.0-design.md decision 4)
         target_s1:  (B, 9) float32 — normalised Stage-1 primary target
         n_sec:      (B,) int64 — true secondary count per step
-        sec_cont:   (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit,
+        sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
                     local_dir, log_mass, charge] per slot (mass/charge
                     normalised iff `sec_phys_normalizer` was given); always
                     computed the same way regardless of
@@ -54,9 +55,13 @@ class StreamingStepsDataset(IterableDataset):
                     downstream under target="physical"
         proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
                     only; zeros when `proc_map` is None)
-        sec_type_idx: (B, K_MAX) int64 — per-slot class index into
+        sec_type_idx: (B, k_max) int64 — per-slot class index into
                     `sec_type_class_map`, for particle_type.target in
                     ("onehot", "embedding"); zeros (unused) otherwise
+
+    `k_max` (constructor arg, default the module constant) should match
+    `stage2_model.k_max` (docs/v0.3.0-design.md §8) — it sets the padded
+    width of `sec_cont`/`sec_type_idx` above.
     """
 
     def __init__(
@@ -71,11 +76,13 @@ class StreamingStepsDataset(IterableDataset):
         shuffle_buffer: int = 65536,
         shuffle: bool = True,
         proc_map: dict[str, int] | None = None,
-        conditioning: str = "embedding",
+        particle_conditioning: str = "embedding",
+        material_conditioning: str = "embedding",
         sec_phys_normalizer: Normalizer | None = None,
         pdg_topn_map: dict[int, int] | None = None,
         mat_topn_map: dict[str, int] | None = None,
         sec_type_class_map: dict | None = None,
+        k_max: int = K_MAX,
     ) -> None:
         self.files = list(files)
         self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
@@ -89,11 +96,13 @@ class StreamingStepsDataset(IterableDataset):
         self.shuffle_buffer = max(shuffle_buffer, batch_size)
         self.shuffle = shuffle
         self.proc_map = proc_map
-        self.conditioning = conditioning
+        self.particle_conditioning = particle_conditioning
+        self.material_conditioning = material_conditioning
         self.sec_phys_normalizer = sec_phys_normalizer
         self.pdg_topn_map = pdg_topn_map
         self.mat_topn_map = mat_topn_map
         self.sec_type_class_map = sec_type_class_map
+        self.k_max = k_max
 
     def __iter__(self):
         worker_info = torch.utils.data.get_worker_info()
@@ -115,7 +124,9 @@ class StreamingStepsDataset(IterableDataset):
         buf_n = 0
 
         for path in files:
-            for chunk in iter_file_chunks(path, offset=self._offsets[path]):
+            for chunk in iter_file_chunks(
+                path, offset=self._offsets[path], k_max=self.k_max
+            ):
                 mask = sorted_membership(chunk["event_id"], self._events_arr)
                 if not mask.any():
                     continue
@@ -140,10 +151,12 @@ class StreamingStepsDataset(IterableDataset):
                     sec_phys_normalizer=self.sec_phys_normalizer,
                     proc_map=self.proc_map,
                     require_secondaries=True,
-                    conditioning=self.conditioning,
+                    particle_conditioning=self.particle_conditioning,
+                    material_conditioning=self.material_conditioning,
                     pdg_topn_map=self.pdg_topn_map,
                     mat_topn_map=self.mat_topn_map,
                     sec_type_class_map=self.sec_type_class_map,
+                    k_max=self.k_max,
                 )
                 buf_cont.append(cond_cont)
                 buf_cat.append(cond_cat)

giant/data/loader.py

@@ -6,6 +6,8 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
 
+from giant.constants import K_MAX
+
 # A manifest is a plain text file listing one parquet path per line, used to
 # name a curated subset of files (e.g. a train/holdout pool) without copying
 # or symlinking the underlying parquet files. Lines are resolved relative to
@@ -113,9 +115,9 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
     return out
 
 
-def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
-    from giant.constants import K_MAX
-
+def _df_to_dict(
+    df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX
+) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
@@ -147,17 +149,19 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     }
 
     if has_sec_lists:
-        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
-        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
+        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
+        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
         d["sec_dir_list"] = _pad_dir_col(
-            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
+            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max
         )
 
     return d
 
 
-def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
-    return _df_to_dict(pd.read_parquet(path), offset=offset)
+def load_steps(
+    path: str | Path, offset: int = 0, k_max: int = K_MAX
+) -> dict[str, np.ndarray]:
+    return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
 
 
 def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
@@ -167,12 +171,16 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
 
 
 def iter_file_chunks(
-    path: str | Path, offset: int = 0
+    path: str | Path, offset: int = 0, k_max: int = K_MAX
 ) -> Iterator[dict[str, np.ndarray]]:
-    """Yield one parquet row-group at a time so a large file never fully loads."""
+    """Yield one parquet row-group at a time so a large file never fully loads.
+
+    `k_max` sets the padded width of the sec_*_list columns (should match
+    `stage2_model.k_max` — see docs/v0.3.0-design.md §8); defaults to the
+    module constant for callers that don't care (e.g. Stage-1-only reads)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)
+        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)
 
 
 _COND_COLS = [

giant/data/setup_cache.py

@@ -96,10 +96,22 @@ def fingerprint_files(files: list[Path]) -> list[list]:
     return out
 
 
-def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
+def normalizer_key(
+    val_fraction: float,
+    seed: int,
+    particle_conditioning: str,
+    material_conditioning: str,
+) -> str:
     # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
-    # spurious cache misses between runs with the "same" val_fraction.
-    return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"
+    # spurious cache misses between runs with the "same" val_fraction. The two
+    # conditioning axes are independent (docs/v0.3.0-design.md §3.1) and both
+    # affect which cond_cont columns are computed for real vs. zero-filled
+    # (giant.data.transforms._physical_cond_columns), so both must be part of
+    # the key or two mixed-axis runs could collide on the same cache entry.
+    return (
+        f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}"
+        f"_mcond={material_conditioning}"
+    )
 
 
 # Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int

giant/data/transforms.py

@@ -2,6 +2,8 @@ import warnings
 
 import numpy as np
 
+from giant.constants import K_MAX
+
 _EPS = 1e-8
 
 # Floor added to each energy fraction before taking log-ratios so the simplex
@@ -719,14 +721,23 @@ def decode_secondaries(
 
 
 def _physical_cond_columns(
-    data: dict[str, np.ndarray], conditioning: str
+    data: dict[str, np.ndarray],
+    particle_conditioning: str,
+    material_conditioning: str,
 ) -> np.ndarray:
     """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
 
-    "embedding"/"onehot" modes zero-fill (cheap, and ConditionEncoder never
-    reads these columns in either mode — so an unfilled giant.materials table
-    can never crash an "embedding"/"onehot"-mode run). "physical" mode
-    computes them for real: particle columns come from
+    The particle and material blocks are gated independently
+    (docs/v0.3.0-design.md §3.1: "configured independently and may mix
+    freely — e.g. material `physical` with particle `embedding`"), so e.g.
+    `particle_conditioning="embedding"` + `material_conditioning="physical"`
+    zero-fills only the particle columns and computes the material ones for
+    real.
+
+    "embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder
+    never reads these columns in either mode — so an unfilled
+    giant.materials table can never crash an "embedding"/"onehot"-mode run).
+    "physical" computes it for real: particle columns come from
     `data["mass"]`/`data["charge"]` when the caller already knows them
     directly (rollout.py, for a track descended from a model-predicted
     secondary — see giant/rollout.py's "no snapping" design), else derived
@@ -736,36 +747,47 @@ def _physical_cond_columns(
     """
     from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
 
-    if conditioning in ("embedding", "onehot"):
-        n = len(next(iter(data.values())))
-        return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32)
-    if conditioning != "physical":
-        raise ValueError(f"unknown conditioning mode {conditioning!r}")
+    n = len(next(iter(data.values())))
 
-    from giant.materials import material_properties_array
-    from giant.particles import particle_phys_array
+    if particle_conditioning == "physical":
+        from giant.particles import particle_phys_array
 
-    if "mass" in data and "charge" in data:
-        mass = np.asarray(data["mass"], dtype=np.float32)
-        charge = np.asarray(data["charge"], dtype=np.float32)
+        if "mass" in data and "charge" in data:
+            mass = np.asarray(data["mass"], dtype=np.float32)
+            charge = np.asarray(data["charge"], dtype=np.float32)
+        else:
+            mass, charge = particle_phys_array(data["pdg"]).T
+        particle_cols = np.column_stack([log_transform(mass), charge])
+    elif particle_conditioning in ("embedding", "onehot"):
+        particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
     else:
-        mass, charge = particle_phys_array(data["pdg"]).T
+        raise ValueError(
+            f"unknown conditioning.particle.type {particle_conditioning!r}"
+        )
 
-    z_eff, a_eff, density, x0, lambda_int = material_properties_array(
-        data["material"]
-    ).T
+    if material_conditioning == "physical":
+        from giant.materials import material_properties_array
 
-    return np.column_stack(
-        [
-            log_transform(mass),
-            charge,
-            z_eff,
-            a_eff,
-            log_transform(density),
-            log_transform(x0),
-            log_transform(lambda_int),
-        ]
-    ).astype(np.float32)
+        z_eff, a_eff, density, x0, lambda_int = material_properties_array(
+            data["material"]
+        ).T
+        material_cols = np.column_stack(
+            [
+                z_eff,
+                a_eff,
+                log_transform(density),
+                log_transform(x0),
+                log_transform(lambda_int),
+            ]
+        )
+    elif material_conditioning in ("embedding", "onehot"):
+        material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
+    else:
+        raise ValueError(
+            f"unknown conditioning.material.type {material_conditioning!r}"
+        )
+
+    return np.column_stack([particle_cols, material_cols]).astype(np.float32)
 
 
 def build_cond_features(
@@ -773,19 +795,24 @@ def build_cond_features(
     pdg_map: dict[int, int],
     mat_map: dict[str, int],
     cond_normalizer: "Normalizer | None" = None,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
     pdg_topn_map: dict[int, int] | None = None,
     mat_topn_map: dict[str, int] | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Build conditioning arrays only — no target, no post-step variables.
 
+    `particle_conditioning`/`material_conditioning` are independent
+    (docs/v0.3.0-design.md §3.1) — e.g. `particle_conditioning="embedding"` +
+    `material_conditioning="physical"` is a valid mix.
+
     `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
     `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
     columns read by `ConditionEncoder`'s `"onehot"` mode
     (docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff
     `pdg_topn_map` given), material topN index at column 3 (iff
     `mat_topn_map` given, after column 2 if both are). Only ever given when
-    `conditioning == "onehot"`; `cond_cat` stays `(N, 2)` otherwise.
+    the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise.
     """
     cond_cont = np.column_stack(
         [
@@ -796,7 +823,10 @@ def build_cond_features(
         ]
     ).astype(np.float32)
     cond_cont = np.column_stack(
-        [cond_cont, _physical_cond_columns(data, conditioning)]
+        [
+            cond_cont,
+            _physical_cond_columns(data, particle_conditioning, material_conditioning),
+        ]
     ).astype(np.float32)
 
     # In "physical" mode cond_cat's first two columns are only a
@@ -807,10 +837,11 @@ def build_cond_features(
     # conditioning signal, so an unmapped value must still raise loudly
     # rather than silently misassign. In "onehot" mode they again go unread
     # (the topN columns below are the real signal), so they're as permissive
-    # as "physical".
-    strict = conditioning == "embedding"
-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict)
+    # as "physical". Each axis's strictness is independent.
+    pdg_strict = particle_conditioning == "embedding"
+    mat_strict = material_conditioning == "embedding"
+    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=pdg_strict)
+    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=mat_strict)
     cat_cols = [pdg_idx, mat_idx]
     if pdg_topn_map is not None:
         cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
@@ -819,34 +850,44 @@ def build_cond_features(
     cond_cat = np.column_stack(cat_cols)
 
     if cond_normalizer is not None:
-        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning)
+        cond_cont = _cond_normalizer_transform(
+            cond_cont, cond_normalizer, particle_conditioning, material_conditioning
+        )
 
     return cond_cont, cond_cat
 
 
 def _cond_normalizer_transform(
-    cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str
+    cond_cont: np.ndarray,
+    cond_normalizer: "Normalizer",
+    particle_conditioning: str,
+    material_conditioning: str,
 ) -> np.ndarray:
     """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.
 
     Checkpoints trained before physical-property conditioning (``COND_DIM``
     8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
     normalizer, fit before ``build_cond_features`` grew the extra physical
-    columns. In "embedding" mode those columns are never read by
+    columns. When NEITHER axis is "physical" those columns are never read by
     ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing
     entries with mean=0/std=1 is a safe no-op that keeps such checkpoints
-    usable under the current, always-``COND_DIM``-wide contract. In
-    "physical" mode the physical columns are load-bearing, so a mismatch
-    there is a real incompatibility, not something to paper over.
+    usable under the current, always-``COND_DIM``-wide contract. If EITHER
+    axis is "physical" its columns are load-bearing, so a mismatch there is a
+    real incompatibility, not something to paper over.
     """
     mean, std = cond_normalizer.mean, cond_normalizer.std
     assert mean is not None and std is not None, "Normalizer not fitted"
     width = cond_cont.shape[-1]
     if mean.shape[-1] < width:
-        if conditioning != "embedding":
+        physical_load_bearing = "physical" in (
+            particle_conditioning,
+            material_conditioning,
+        )
+        if physical_load_bearing:
             raise ValueError(
                 f"cond normalizer has {mean.shape[-1]} columns, expected "
-                f"{width}, and conditioning={conditioning!r} reads the "
+                f"{width}, and particle_conditioning={particle_conditioning!r}/"
+                f"material_conditioning={material_conditioning!r} reads the "
                 "physical columns directly — this checkpoint predates "
                 "physical-property conditioning and can't be safely padded; "
                 "retrain it under the current code."
@@ -867,11 +908,13 @@ def build_features(
     fit: bool = False,
     proc_map: dict[str, int] | None = None,
     require_secondaries: bool = False,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
     sec_phys_only: bool = False,
     pdg_topn_map: dict[int, int] | None = None,
     mat_topn_map: dict[str, int] | None = None,
     sec_type_class_map: dict | None = None,
+    k_max: int = K_MAX,
 ) -> tuple[
     np.ndarray,
     np.ndarray,
@@ -922,8 +965,14 @@ def build_features(
         top-N-plus-other map's `class_map` for `target = "onehot"`, or the
         dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
         `None` for `target = "physical"`.
+
+    k_max: should match `stage2_model.k_max` (docs/v0.3.0-design.md §8) —
+        overridden internally by `data["sec_E_list"]`'s own padded width when
+        present (the loader already padded it to some k_max; that width is
+        authoritative), so this only actually matters when secondary list
+        columns are absent (Stage-1-only reads, or a pre-secondary-join
+        file), where it sets `sec_cont`/`sec_type_idx`'s zero-filled width.
     """
-    from giant.constants import K_MAX
 
     post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
     travel_dir_local = local_frame_rotation(
@@ -953,7 +1002,10 @@ def build_features(
         ]
     ).astype(np.float32)  # (N, COND_DIM_BASE=8)
     cond_cont = np.column_stack(
-        [cond_cont, _physical_cond_columns(data, conditioning)]
+        [
+            cond_cont,
+            _physical_cond_columns(data, particle_conditioning, material_conditioning),
+        ]
     ).astype(np.float32)  # (N, COND_DIM=15)
 
     pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
@@ -968,20 +1020,27 @@ def build_features(
     n_sec_raw = data["n_sec"].astype(
         np.int64
     )  # (N,) unclamped, for the valid-slot mask
-    # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes
-    # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already
-    # applied to sec_cont by the loader's list padding. Without this, a rare
-    # high-multiplicity step (real data goes up to ~37) hands cross_entropy
-    # an out-of-range target and CUDA asserts.
-    n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64)  # (N,)
 
     # Secondary continuous targets
     sec_E_list = data.get("sec_E_list")
     sec_dir_list = data.get("sec_dir_list")
     sec_pdg_list = data.get("sec_pdg_list")
+    if sec_E_list is not None:
+        # The loader already padded sec_*_list to some k_max (see
+        # giant.data.loader.iter_file_chunks); that padded width is
+        # authoritative over whatever this call happened to pass in, so the
+        # two can never drift apart (docs/v0.3.0-design.md §8).
+        k_max = sec_E_list.shape[1]
+
+    # Clamp the classification label to k_max: the head only has k_max+1
+    # classes (0..k_max), and truncating here mirrors the k_max-slot
+    # truncation already applied to sec_cont by the loader's list padding.
+    # Without this, a rare high-multiplicity step (real data goes up to ~37)
+    # hands cross_entropy an out-of-range target and CUDA asserts.
+    n_sec = np.minimum(n_sec_raw, k_max).astype(np.int64)  # (N,)
 
     if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
-        sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None]  # (N, K_MAX)
+        sec_valid = np.arange(k_max)[None, :] < n_sec_raw[:, None]  # (N, k_max)
         sec_cont = encode_secondaries(
             sec_E_list,
             sec_dir_list,
@@ -990,11 +1049,11 @@ def build_features(
             data["pre_dir"],
             sec_pdg_list=sec_pdg_list,
             phys_only=sec_phys_only,
-        )  # (N, K_MAX, 6)
+        )  # (N, k_max, 6)
         sec_type_idx = (
             encode_secondary_type_idx(sec_pdg_list, sec_valid, sec_type_class_map)
             if sec_type_class_map is not None
-            else np.zeros((len(n_sec), K_MAX), dtype=np.int64)
+            else np.zeros((len(n_sec), k_max), dtype=np.int64)
         )
     else:
         # Guard against silently training Stage 2 on zeroed targets: if any step
@@ -1017,8 +1076,8 @@ def build_features(
                 "require_secondaries=False for Stage-1-only use."
             )
         N = len(n_sec)
-        sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
-        sec_type_idx = np.zeros((N, K_MAX), dtype=np.int64)
+        sec_cont = np.zeros((N, k_max, 6), dtype=np.float32)
+        sec_type_idx = np.zeros((N, k_max), dtype=np.int64)
 
     if fit:
         cond_normalizer = Normalizer().fit(cond_cont)

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)
 # ---------------------------------------------------------------------------
@@ -863,7 +990,12 @@ class Stage1Model(nn.Module):
     (docs/v0.3.0-design.md §2) moves it to stage 2, except for a migrated
     v0.2 checkpoint (`n_sec_head_k_max` given), where it stays attached here
     since that's where its weights live and what conditioning it was trained
-    against (see `_migrate_legacy_model_config`)."""
+    against (see `_migrate_legacy_model_config`).
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — `conditioning.share_stages = true` (§3.1): `build_models`
+    constructs one shared instance and passes it to both stages, halving the
+    conditioning parameter count and forcing a common representation."""
 
     def __init__(
         self,
@@ -881,12 +1013,17 @@ class Stage1Model(nn.Module):
         noise_dim: int = 64,
         router: Router | None = None,
         n_sec_head_k_max: int | None = None,
+        cond_enc: ConditionEncoder | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
         self.noise_dim = noise_dim
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(
+                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+            )
         )
         has_time = generator in ("flow", "ddpm")
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
@@ -956,6 +1093,9 @@ class Stage2OneShot(nn.Module):
     Under `generator == "wgan"` the type slice stays folded into `sec_dim`
     (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
     unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
     """
 
     def __init__(
@@ -978,6 +1118,7 @@ class Stage2OneShot(nn.Module):
         router: Router | None = None,
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
+        cond_enc: ConditionEncoder | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
@@ -985,8 +1126,12 @@ class Stage2OneShot(nn.Module):
         self.k_max = k_max
         self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
         self.type_dim = stage2_type_dim(self.particle_type_cfg, particle_cfg["emb_dim"])
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(
+                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+            )
         )
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.fuse = nn.Sequential(
@@ -1082,8 +1227,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.
@@ -1100,6 +1245,9 @@ class Stage2Autoregressive(nn.Module):
     on token position; `_token_cond` additionally fuses in the history
     encoding and two running scalars (remaining energy-budget fraction,
     normalized slot index), and feeds `forward`/`predict_type`/the trunk.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
     """
 
     def __init__(
@@ -1122,13 +1270,17 @@ 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,
+        cond_enc: ConditionEncoder | None = None,
     ) -> 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
@@ -1136,8 +1288,12 @@ class Stage2Autoregressive(nn.Module):
         emb_dim = particle_cfg["emb_dim"]
         self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)
 
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(
+                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+            )
         )
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
@@ -1149,7 +1305,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 +1368,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 +1421,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 +1432,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 +1467,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 +1485,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)
@@ -1387,6 +1588,20 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
     k_max = m.get("k_max", K_MAX)
     noise_dim = m.get("noise_dim", 64)
     router_cfg = dict(m.get("router") or {})
+    expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
+    expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
+    if expert_hidden_dim or expert_n_blocks:
+        raise ValueError(
+            "this checkpoint's model_config.router sets expert_hidden_dim/"
+            f"expert_n_blocks to a non-default value "
+            f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed "
+            "per-expert sizing (experts always inherit the stage's "
+            "hidden_dim/n_res_blocks), so this checkpoint's routed experts "
+            "have a different width/depth than the monolith — silently "
+            "dropping these keys would resize the experts instead of "
+            "refusing (docs/v0.3.0-design.md §4.2). This checkpoint can "
+            "only be loaded by v0.2 code."
+        )
     router_cfg.setdefault("enabled", False)
 
     return {
@@ -1482,6 +1697,13 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     instance rather than building a second, independently-parameterized one
     (v0.2's actual — probably accidental — behaviour: two routers built from
     one config with no semantic relationship between them).
+
+    `conditioning.share_stages = true` (§3.1) builds one `ConditionEncoder`
+    instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
+    `Stage2Autoregressive`'s `cond_enc` param), instead of each stage
+    building its own — halving the conditioning parameter count and forcing a
+    common representation. `false` (default) keeps v0.2 behaviour:
+    independent instances with identical config but independent weights.
     """
     cfg = (
         model_config
@@ -1491,17 +1713,17 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     pdg_vocab = cfg["pdg_vocab"]
     mat_vocab = cfg["mat_vocab"]
     conditioning = cfg["conditioning"]
-    if conditioning.get("share_stages"):
-        raise NotImplementedError(
-            "conditioning.share_stages = true is not implemented yet — each "
-            "stage always builds its own ConditionEncoder for now"
-        )
     particle_cfg = conditioning["particle"]
     material_cfg = conditioning["material"]
     particle_conditioning = particle_cfg["type"]
     s1cfg = cfg["stage1_model"]
     s2cfg = cfg["stage2_model"]
     cond_out_dim = conditioning.get("out_dim", 128)
+    shared_cond_enc: ConditionEncoder | None = None
+    if conditioning.get("share_stages"):
+        shared_cond_enc = ConditionEncoder(
+            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        )
 
     result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
 
@@ -1532,6 +1754,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             noise_dim=(s1cfg.get("wgan") or {}).get("noise_dim", 64),
             router=stage1_router,
             n_sec_head_k_max=n_sec_head_k_max,
+            cond_enc=shared_cond_enc,
         )
 
     if s2cfg.get("active", True):
@@ -1571,6 +1794,9 @@ 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),
+                cond_enc=shared_cond_enc,
             )
         else:
             sec_dim = stage2_trunk_sec_dim(
@@ -1594,6 +1820,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 router=stage2_router,
                 build_n_sec_head=legacy_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
+                cond_enc=shared_cond_enc,
             )
 
     return result

giant/model/schedule.py

@@ -108,12 +108,13 @@ def flow_matching_loss_secondary(
     averaged over their own width first and then combined with equal weight
     — this stays correct if type_dim/CONT_SLOT_DIM change.
     """
-    from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM
+    from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
 
     if type_dim is None:
         type_dim = PARTICLE_PHYS_DIM
 
     B = x1.size(0)
+    k_max = sec_mask.size(1)
     t = torch.rand(B, device=x1.device)
     x0 = torch.randn_like(x1)
     x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
@@ -121,8 +122,8 @@ def flow_matching_loss_secondary(
     v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
 
     slot_dim = CONT_SLOT_DIM + type_dim
-    err = ((v_t - u_t) ** 2).view(B, K_MAX, slot_dim)
-    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, K_MAX)
+    err = ((v_t - u_t) ** 2).view(B, k_max, slot_dim)
+    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, k_max)
 
     mask = sec_mask.float()
     denom = mask.sum().clamp(min=1)

giant/pipeline.py

@@ -9,7 +9,6 @@ from torch.utils.data import DataLoader
 from giant import config
 from giant.constants import (
     COND_DIM,
-    K_MAX,
     PARTICLE_PHYS_DIM,
     X_DIM,
 )
@@ -106,11 +105,9 @@ def run_setup_stage(
     Stage-2 normalizers.
 
     `cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/
-    `stage2_model`), already passed through `giant.config.validate_config` —
-    in particular, `conditioning.particle.type == conditioning.material.type`
-    is assumed (the data pipeline doesn't support mixed conditioning types
-    yet, see `validate_config`), so a single shared conditioning string
-    drives `giant.data.transforms.build_features`.
+    `stage2_model`), already passed through `giant.config.validate_config`.
+    `conditioning.particle.type` and `conditioning.material.type` are
+    independent (docs/v0.3.0-design.md §3.1) and may differ.
 
     Reads from and writes to the `giant.data.setup_cache` sidecar when
     `cache_setup` is set (`rebuild_setup_cache` ignores — but still
@@ -118,7 +115,9 @@ def run_setup_stage(
     is mutated in place: an active `EnergyRouter` (`router.type == "energy"`)
     gets its `centers_init` seeded from real data quantiles here.
     """
-    conditioning = cfg["conditioning"]["particle"]["type"]
+    particle_conditioning = cfg["conditioning"]["particle"]["type"]
+    material_conditioning = cfg["conditioning"]["material"]["type"]
+    k_max = cfg["stage2_model"]["k_max"]
     stage1_router = cfg["stage1_model"].get("router") or {}
     stage2_router = cfg["stage2_model"].get("router") or {}
 
@@ -255,7 +254,9 @@ def run_setup_stage(
         for r in (stage1_router, stage2_router)
     )
     energy_idx = 3
-    norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning)
+    norm_key = setup_cache.normalizer_key(
+        val_fraction, seed, particle_conditioning, material_conditioning
+    )
     entry = cache.normalizers.get(norm_key) if cache is not None else None
 
     if entry is not None:
@@ -284,7 +285,7 @@ def run_setup_stage(
             _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
         )
         for i, path in enumerate(files):
-            for chunk in iter_file_chunks(path, offset=event_id_offset(i)):
+            for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
                 mask = sorted_membership(chunk["event_id"], events_arr)
                 if not mask.any():
                     continue
@@ -296,15 +297,17 @@ def run_setup_stage(
                         mat_map,
                         proc_map=proc_map,
                         require_secondaries=True,
-                        conditioning=conditioning,
+                        particle_conditioning=particle_conditioning,
+                        material_conditioning=material_conditioning,
                         sec_phys_only=True,
+                        k_max=k_max,
                     )
                 )
                 cond_acc.update(cond_cont)
                 tgt_acc.update(target_s1)
                 if energy_sampler is not None:
                     energy_sampler.update(cond_cont[:, energy_idx])
-                sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
+                sec_valid = np.arange(sec_cont.shape[1])[None, :] < n_sec[:, None]
                 sec_phys = sec_cont[:, :, 4:6][sec_valid]
                 if len(sec_phys) > 0:
                     sec_phys_acc.update(sec_phys)
@@ -375,7 +378,9 @@ def run_train_job(
         )
 
     config.validate_config(cfg)
-    conditioning = cfg["conditioning"]["particle"]["type"]
+    particle_conditioning = cfg["conditioning"]["particle"]["type"]
+    material_conditioning = cfg["conditioning"]["material"]["type"]
+    k_max = cfg["stage2_model"]["k_max"]
     setup = run_setup_stage(
         data,
         val_fraction=t["val_fraction"],
@@ -399,18 +404,20 @@ def run_train_job(
     )
 
     # cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are
-    # present only under conditioning="onehot" — validate_config enforces
-    # conditioning.particle.type == conditioning.material.type, so both
-    # axes' topN columns are always present or absent together. run_setup_stage
-    # builds both maps whenever conditioning=="onehot" (see its own
-    # particle_cfg["type"] == "onehot" check), so they're guaranteed non-None
-    # here — asserted, not just assumed, so a future wiring bug fails loudly
-    # instead of silently dropping the onehot columns.
+    # present per-axis, independently, under that axis's own
+    # conditioning.{particle,material}.type == "onehot" (§3.1: the two axes
+    # may mix freely). run_setup_stage builds each map whenever its own axis
+    # is "onehot" (see its own particle_cfg["type"]/material_cfg["type"]
+    # checks), so they're guaranteed non-None here — asserted, not just
+    # assumed, so a future wiring bug fails loudly instead of silently
+    # dropping the onehot columns.
     cond_pdg_topn = None
     cond_mat_topn = None
-    if conditioning == "onehot":
-        assert setup.pdg_topn_map is not None and setup.mat_topn_map is not None
+    if particle_conditioning == "onehot":
+        assert setup.pdg_topn_map is not None
         cond_pdg_topn = setup.pdg_topn_map.class_map
+    if material_conditioning == "onehot":
+        assert setup.mat_topn_map is not None
         cond_mat_topn = setup.mat_topn_map.class_map
 
     # The secondary type-index map depends on stage2_model.particle_type.target,
@@ -441,11 +448,13 @@ def run_train_job(
         shuffle_buffer=shuffle_buffer,
         shuffle=True,
         proc_map=proc_map,
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         sec_phys_normalizer=sec_phys_norm,
         pdg_topn_map=cond_pdg_topn,
         mat_topn_map=cond_mat_topn,
         sec_type_class_map=sec_type_class_map,
+        k_max=k_max,
     )
     val_ds = StreamingStepsDataset(
         files=files,
@@ -457,11 +466,13 @@ def run_train_job(
         batch_size=t["batch_size"],
         shuffle=False,
         proc_map=proc_map,
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         sec_phys_normalizer=sec_phys_norm,
         pdg_topn_map=cond_pdg_topn,
         mat_topn_map=cond_mat_topn,
         sec_type_class_map=sec_type_class_map,
+        k_max=k_max,
     )
 
     pin = device.type == "cuda"

giant/rollout.py

@@ -352,7 +352,7 @@ def make_seed_frontier(
     pre_pos: np.ndarray,
     pre_E: np.ndarray,
     pre_dir: np.ndarray,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
 ) -> tuple[dict[str, np.ndarray], dict[int, int]]:
     """Build the initial frontier from primary entry states.
 
@@ -372,17 +372,17 @@ def make_seed_frontier(
     dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)
 
     pdg_arr = np.asarray(pdg, dtype=np.int64)
-    if conditioning == "physical":
+    if particle_conditioning == "physical":
         # Real primaries always have a genuine ground-truth PDG code, looked
         # up once here and carried forward unchanged for the track's lifetime
         # (its species never changes mid-track) — same lifecycle as "pdg"
         # itself.
         mass, charge = particle_phys_array(pdg_arr).T
     else:
-        # "embedding" mode never reads mass/charge (see
+        # "embedding"/"onehot" never read mass/charge (see
         # _physical_cond_columns), so resolving them here would only risk
-        # crashing an embedding-mode rollout on a PDG code giant.particles
-        # can't resolve, for a value that's never used.
+        # crashing a rollout on a PDG code giant.particles can't resolve, for
+        # a value that's never used.
         mass = np.zeros(n, dtype=np.float64)
         charge = np.zeros(n, dtype=np.float64)
 
@@ -457,8 +457,10 @@ def rollout(
     max_tracks_per_event: int | None = None,
     escape_threshold: float | None = None,
     on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
     pdg_topn_map: "TopNMap | None" = None,
+    mat_topn_map: "TopNMap | None" = None,
     other_policy: str = "sample",
     seed: int | None = None,
     stage1_ddpm_steps: int = 1000,
@@ -485,8 +487,13 @@ def rollout(
     autoregressive) is inferred from `sec_decoder`'s own class — see
     `sample_stage1`/`sample_stage2` (giant.sample).
 
-    `pdg_topn_map`/`other_policy` are only read under
-    `stage2_model.particle_type.target = "onehot"` (§3.3); `seed` seeds the
+    `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen
+    to share `pdg_topn_map` (docs/v0.3.0-design.md §8 — one PDG map, not
+    two): they're required whenever `particle_conditioning`/
+    `material_conditioning` is `"onehot"` (feeds `build_cond_features`'s
+    extra `cond_cat` top-N columns — §3.1), and `pdg_topn_map`/`other_policy`
+    are additionally read under `stage2_model.particle_type.target =
+    "onehot"` (§3.3, secondary-species decode). `seed` seeds the
     `other_policy = "sample"` draw only (torch/numpy sampling itself is
     seeded by the caller, same as today).
 
@@ -494,6 +501,16 @@ def rollout(
     diagnostic across the whole run — see `L1DistCollector`. Only populated
     under `particle_type.target = "embedding"`; a no-op otherwise.
     """
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        raise RuntimeError(
+            "conditioning.particle.type='onehot' rollout needs pdg_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
+        )
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        raise RuntimeError(
+            "conditioning.material.type='onehot' rollout needs mat_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['mat_topn_map']"
+        )
     device = device or torch.device("cpu")
     stage1_model.eval()
     sec_decoder.eval()
@@ -507,7 +524,7 @@ def rollout(
         seeds["pre_pos"],
         seeds["pre_E"],
         seeds["pre_dir"],
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
     )
     rec = _Recorder(sink=on_chunk)
 
@@ -534,8 +551,10 @@ def rollout(
                     steps,
                     device,
                     max_tracks_per_event,
-                    conditioning,
+                    particle_conditioning,
+                    material_conditioning,
                     pdg_topn_map,
+                    mat_topn_map,
                     other_policy,
                     rng,
                     stage1_ddpm_steps,
@@ -570,8 +589,10 @@ def _step_chunk(
     steps,
     device,
     max_tracks_per_event,
-    conditioning,
+    particle_conditioning,
+    material_conditioning,
     pdg_topn_map,
+    mat_topn_map,
     other_policy,
     rng,
     stage1_ddpm_steps,
@@ -587,7 +608,7 @@ def _step_chunk(
     tr["_material"] = material
     tr["_layer_id"] = layer_id
 
-    if conditioning == "physical":
+    if particle_conditioning == "physical":
         # Under physical-property conditioning, mass/charge (already resolved
         # on every track — see the cond_dict comment below) drive the model,
         # not a training-vocab PDG embedding — build_cond_features passes
@@ -655,7 +676,18 @@ def _step_chunk(
         "charge": tr["charge"],
     }
     cond_cont, cond_cat = build_cond_features(
-        cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning
+        cond_dict,
+        pdg_map,
+        mat_map,
+        cond_norm,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
+        pdg_topn_map=pdg_topn_map.class_map
+        if particle_conditioning == "onehot"
+        else None,
+        mat_topn_map=mat_topn_map.class_map
+        if material_conditioning == "onehot"
+        else None,
     )
     cc = torch.from_numpy(cond_cont).float().to(device)
     ck = torch.from_numpy(cond_cat).long().to(device)

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

giant/train.py

@@ -4,7 +4,6 @@ import math
 import os
 import signal
 import time
-import warnings
 from pathlib import Path
 from types import FrameType
 from typing import Callable
@@ -16,7 +15,7 @@ import torch.optim as optim
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM
+from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_to_json
 from giant.model.network import Router, stage2_type_dim
@@ -27,6 +26,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 +246,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,17 +475,28 @@ 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(
-                f"stage2_model.generator={generator!r} is not implemented for "
-                "stage 2 (only 'flow' and 'wgan' — 'ddpm' has no stage-2 "
-                "secondary-decoder loss yet, see docs/v0.3.0-design.md)"
+                f"stage2_model.generator={generator!r} is accepted by the "
+                "schema but not implemented in v0.3.0 for stage 2 (only "
+                "'flow' and 'wgan' have a stage-2 secondary-decoder loss — "
+                "see docs/v0.3.0-design.md §11.2)"
             )
         self.name = name
         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,
@@ -523,18 +653,38 @@ class FlowDDPMStageTrainer(StageTrainer):
             proc_idx,
             sec_type_idx,
         ) = _batch_to_device(batch, device)
-        sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
+        sec_mask = torch.arange(sec_cont.size(1), device=device).unsqueeze(
+            0
+        ) < n_sec.unsqueeze(1)
         stage1_ctx = x1_s1.detach()
 
         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 +763,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 +870,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)
@@ -798,8 +959,9 @@ class WGANStageTrainer(StageTrainer):
                 self.particle_type_cfg, self.particle_type_emb_dim
             )
             slot_width = CONT_SLOT_DIM + type_dim
+            k_max = sec_cont.size(1)
 
-            sec_mask = torch.arange(K_MAX, device=device).unsqueeze(
+            sec_mask = torch.arange(k_max, device=device).unsqueeze(
                 0
             ) < n_sec.unsqueeze(1)
             mask = (
@@ -810,12 +972,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(
@@ -828,7 +1005,7 @@ class WGANStageTrainer(StageTrainer):
                     ).reshape(B, -1)
                     * mask
                 )
-                z = torch.randn(B, K_MAX, self.model.noise_dim, device=device)
+                z = torch.randn(B, k_max, self.model.noise_dim, device=device)
                 fake_raw = self.model(
                     z,
                     cond_cont,
@@ -868,7 +1045,7 @@ class WGANStageTrainer(StageTrainer):
                     self.type_gumbel_tau_end,
                 )
                 fake_raw = _relax_onehot_type_slice(
-                    fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe
+                    fake_raw, k_max, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe
                 )
             fake = fake_raw * mask
 
@@ -987,39 +1164,14 @@ class WGANStageTrainer(StageTrainer):
             group["lr"] = resumed_critic_lr
 
 
-_WARNED_MARGINAL_VALIDATION_BROKEN = False
-
-
 def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs):
-    """`giant.sample` itself has a per-stage generator/n_sec-ownership
-    update (v0.3.0 step 6 — `sample_stage1`/`sample_stage2`/`resolve_n_sec`,
-    docs/v0.3.0-design.md §10); `giant/validate.py` hasn't been updated to
-    call it yet, and also still unpacks the val batch as a 6-tuple, which no
-    longer matches `StreamingStepsDataset`'s 7-tuple shape (the `sec_type_idx`
-    column step 4 added). Both are `giant/validate.py`-side gaps, not
-    `giant/sample.py`'s — tracked, not yet scheduled. Until fixed, this
-    degrades gracefully with a one-time warning instead of crashing the
-    whole training run.
-    """
-    global _WARNED_MARGINAL_VALIDATION_BROKEN
+    """Runs `validate_marginals` on `trainer`'s sampling model (EMA model if
+    present, else the raw model). `validate_marginals` itself dispatches
+    through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`
+    (docs/v0.3.0-design.md §10), so this is generator- and
+    one-shot-vs-autoregressive-agnostic."""
     model = trainer.sampling_model()
-    try:
-        return validate_marginals(model, val_loader, device=device, **kwargs)
-    except Exception as exc:  # noqa: BLE001 — see docstring: any failure here
-        # is expected until giant/validate.py is updated for the v0.3.0
-        # per-stage dispatch + 7-tuple batch shape.
-        if not _WARNED_MARGINAL_VALIDATION_BROKEN:
-            warnings.warn(
-                "marginal validation unavailable this run "
-                f"({type(exc).__name__}: {exc}) — giant/validate.py hasn't "
-                "been updated for v0.3.0's per-stage sample dispatch / "
-                "7-tuple val batch shape yet; val_marginal_kl stays NaN, and "
-                "wgan best-checkpoint selection falls back to the "
-                "Wasserstein-distance magnitude.",
-                stacklevel=2,
-            )
-            _WARNED_MARGINAL_VALIDATION_BROKEN = True
-        return None
+    return validate_marginals(model, val_loader, device=device, **kwargs)
 
 
 def _build_stage_trainers(
@@ -1051,16 +1203,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 +1241,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 +1271,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
 
@@ -1507,7 +1666,6 @@ def train(
                         stage1_tr,
                         val_loader,
                         device,
-                        mode=cfg["stage1_model"]["generator"],
                         sec_decoder=trainers["stage2"].sampling_model()
                         if "stage2" in trainers
                         else None,
@@ -1535,11 +1693,11 @@ def train(
                         stage1_tr,
                         val_loader,
                         device,
-                        mode=cfg["stage1_model"]["generator"],
-                        schedule=stage1_tr.ddpm_schedule
-                        if isinstance(stage1_tr, FlowDDPMStageTrainer)
-                        else None,
                         steps=validate_steps,
+                        ddpm_steps=stage1_tr.ddpm_schedule.T
+                        if isinstance(stage1_tr, FlowDDPMStageTrainer)
+                        and stage1_tr.ddpm_schedule is not None
+                        else 1000,
                         sec_decoder=trainers["stage2"].sampling_model()
                         if "stage2" in trainers
                         else None,

giant/validate.py

@@ -2,23 +2,12 @@ import numpy as np
 import torch
 from torch.utils.data import DataLoader
 
-from giant.constants import K_MAX, LOCAL_TARGET_NAMES
-from giant.sample import (
-    sample_flow,
-    sample_ddpm,
-    sample_ddim,
-    sample_secondaries,
-    sample_wgan,
-    sample_secondaries_wgan,
-)
+from giant.constants import LOCAL_TARGET_NAMES
+from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 
 _SEC_PHYS_NAMES = ["log_mass", "charge"]
 
 
-def _kw(steps: int | None) -> dict[str, int]:
-    return {} if steps is None else {"steps": steps}
-
-
 def _histogram_kl(
     p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
 ) -> float:
@@ -44,15 +33,44 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
     return counts / total if total > 0 else counts
 
 
+def _categorical_kl(
+    real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8
+) -> float:
+    """KL(P_real || Q_gen) between two class-index samples over `n_classes`
+    categories, estimated from bincount fractions. NaN if either side has no
+    valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
+    if len(real_idx) == 0 or len(gen_idx) == 0:
+        return float("nan")
+    p = _bincount_frac(real_idx, n_classes) + eps
+    q = _bincount_frac(gen_idx, n_classes) + eps
+    p /= p.sum()
+    q /= q.sum()
+    return float(np.sum(p * np.log(p / q)))
+
+
+def _embedding_nearest_class(
+    vectors: torch.Tensor, emb_weight: torch.Tensor
+) -> np.ndarray:
+    """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
+    (vocab, emb_dim) — same computation as
+    `giant.particles.decode_embedding_nearest`, but returning the raw class
+    index instead of a decoded PDG code: validate.py only needs a
+    real-vs-generated class-distribution comparison, not a rollout-usable
+    identity, so there's no need for the pdg_map inversion here."""
+    flat = vectors.reshape(-1, vectors.size(-1))
+    dist = (flat.unsqueeze(1) - emb_weight.detach().unsqueeze(0)).abs().sum(-1)
+    nearest = dist.argmin(dim=1)
+    return nearest.reshape(vectors.shape[:-1]).cpu().numpy()
+
+
 def validate_marginals(
-    model: torch.nn.Module,
+    stage1_model: torch.nn.Module,
     val_loader: DataLoader,
-    mode: str = "flow",
-    schedule=None,
     device: torch.device | None = None,
     n_batches: int | None = None,
     kl_bins: int = 50,
-    steps: int | None = None,
+    steps: int = 10,
+    ddpm_steps: int = 1000,
     sec_decoder: torch.nn.Module | None = None,
 ) -> dict[str, np.ndarray | float]:
     """Compare per-dimension marginals of generated vs. real steps.
@@ -61,51 +79,61 @@ def validate_marginals(
     normalised space. `kl_divergence[j]` is KL(real || generated) for
     dimension j, estimated from a shared histogram over both samples.
 
-    `steps` overrides the number of sampler steps (flow ODE steps or DDIM
-    substeps); `None` keeps each sampler's own default. Unused in "ddpm"
-    mode, which always runs the full schedule.
+    Stage 1 is sampled via `giant.sample.sample_stage1`, which dispatches on
+    `stage1_model.generator_kind` — `steps`/`ddpm_steps` are forwarded but
+    only one of them is actually read, depending on that dispatch.
 
-    When `sec_decoder` is given, also validates Stage 2: n_sec distribution
-    (+ classification accuracy), predicted secondary physical-identity
-    (log_mass, charge) marginals, and per-slot energy-fraction marginals —
-    restricted to each side's own valid slots (real: `n_sec`; generated: the
-    Stage-1 head's argmax), since the two need not agree on how many slots
-    are valid. Compared directly in normalised space (no denormalising —
-    KL estimated from a shared per-sample histogram is invariant to a shared
-    affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred",
-    "n_sec_accuracy", "phys_real", "phys_generated", "phys_kl",
-    "energy_fraction_kl"} to the returned dict.
+    When `sec_decoder` is given, also validates Stage 2 via
+    `giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
+    one-shot-vs-autoregressive-agnostic, docs/v0.3.0-design.md §10): n_sec
+    distribution (+ classification accuracy), per-slot energy-fraction
+    marginals, and a particle-type marginal whose shape depends on
+    `sec_decoder.particle_type_cfg["target"]` — restricted to each side's own
+    valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since
+    the two need not agree on how many slots are valid. Adds {"n_sec_real",
+    "n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under
+    `target = "physical"`, {"phys_real", "phys_generated", "phys_kl"}
+    (continuous log_mass/charge marginals — v0.2 behaviour), or under
+    `target` in `("onehot", "embedding")`, {"type_class_real",
+    "type_class_gen", "type_class_kl"} (categorical class-index marginal:
+    argmax for "onehot", L1-nearest conditioning-embedding row for
+    "embedding" — see `_embedding_nearest_class`). Compared directly in
+    normalised space (no denormalising — KL estimated from a shared
+    per-sample histogram/bincount is invariant to a shared affine rescaling
+    of both sides).
     """
     if device is None:
-        device = next(model.parameters()).device
-    model.eval()
+        device = next(stage1_model.parameters()).device
+    stage1_model.eval()
     if sec_decoder is not None:
         sec_decoder.eval()
 
+    k_max = sec_decoder.k_max if sec_decoder is not None else 0
+    target = (
+        sec_decoder.particle_type_cfg.get("target", "physical")
+        if sec_decoder is not None
+        else "physical"
+    )
+
     all_real, all_gen = [], []
     all_n_sec_real, all_n_sec_pred = [], []
     all_phys_real, all_phys_gen = [], []
-    all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
-    all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
+    all_type_class_real, all_type_class_gen = [], []
+    all_frac_real: list[list[np.ndarray]] = [[] for _ in range(k_max)]
+    all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(k_max)]
 
     for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
-        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx).
-        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch
+        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
+        # sec_type_idx).
+        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch
         cond_cont = cond_cont.to(device)
         cond_cat = cond_cat.to(device)
 
-        if mode == "flow":
-            gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
-        elif mode == "ddpm":
-            gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
-        elif mode == "wgan":
-            gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat)
-        else:
-            gen, n_sec_pred = sample_ddim(
-                model, cond_cont, cond_cat, schedule, **_kw(steps)
-            )
+        gen, n_sec_pred = sample_stage1(
+            stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps
+        )
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())
@@ -113,37 +141,41 @@ def validate_marginals(
         if sec_decoder is None:
             continue
 
+        n_sec_pred = resolve_n_sec(
+            stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
+        )
         n_sec_pred_np = n_sec_pred.cpu().numpy()
         n_sec_np = n_sec.numpy()
         all_n_sec_real.append(n_sec_np)
         all_n_sec_pred.append(n_sec_pred_np)
 
-        real_valid = np.arange(K_MAX)[None, :] < n_sec_np[:, None]  # (B, K_MAX)
+        real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None]  # (B, k_max)
         real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
-        real_phys = sec_cont[:, :, 4:6].numpy()  # (B, K_MAX, 2) [log_mass, charge]
 
-        if mode == "wgan":
-            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan(
-                sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
-            )
-        else:
-            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
-                sec_decoder,
-                cond_cont,
-                cond_cat,
-                gen,
-                n_sec_pred,
-                steps=steps if steps is not None else 10,
-            )
+        sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
+            sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
+        )
         gen_frac = 1.0 / (
             1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
         )
-        gen_phys = sec_phys_pred.cpu().numpy()
         gen_valid = sec_valid_pred.cpu().numpy()
 
-        all_phys_real.append(real_phys[real_valid])
-        all_phys_gen.append(gen_phys[gen_valid])
-        for j in range(K_MAX):
+        if target == "physical":
+            real_phys = sec_cont[:, :, 4:6].numpy()  # (B, k_max, 2) [log_mass, charge]
+            gen_phys = sec_type_pred.cpu().numpy()
+            all_phys_real.append(real_phys[real_valid])
+            all_phys_gen.append(gen_phys[gen_valid])
+        else:
+            sec_type_idx_np = sec_type_idx.numpy()
+            all_type_class_real.append(sec_type_idx_np[real_valid])
+            if target == "onehot":
+                gen_class = sec_type_pred.argmax(dim=-1).cpu().numpy()
+            else:  # "embedding"
+                emb_weight = sec_decoder.cond_enc.pdg_emb.weight
+                gen_class = _embedding_nearest_class(sec_type_pred, emb_weight)
+            all_type_class_gen.append(gen_class[gen_valid])
+
+        for j in range(k_max):
             all_frac_real[j].append(real_frac[real_valid[:, j], j])
             all_frac_gen[j].append(gen_frac[gen_valid[:, j], j])
 
@@ -181,20 +213,8 @@ def validate_marginals(
     n_sec_real = np.concatenate(all_n_sec_real, axis=0)
     n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0)
     n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
-    phys_real = np.concatenate(all_phys_real, axis=0)  # (M, 2)
-    phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
 
-    if len(phys_real) > 0 and len(phys_gen) > 0:
-        phys_kl = np.array(
-            [
-                _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
-                for j in range(2)
-            ]
-        )
-    else:
-        phys_kl = np.full(2, np.nan)
-
-    energy_fraction_kl = np.full(K_MAX, np.nan)
+    energy_fraction_kl = np.full(k_max, np.nan)
     print(
         f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  "
         f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
@@ -208,26 +228,12 @@ def validate_marginals(
     for v in range(max_n_sec):
         print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}")
 
-    print(
-        f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
-        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
-    )
-    print("-" * 68)
-    for j, name in enumerate(_SEC_PHYS_NAMES):
-        r, g = phys_real[:, j], phys_gen[:, j]
-        if len(r) == 0 or len(g) == 0:
-            continue
-        print(
-            f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
-            f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
-        )
-
     print(
         f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
         f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
     )
     print("-" * 90)
-    for j in range(K_MAX):
+    for j in range(k_max):
         r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([])
         g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([])
         if len(r) == 0 or len(g) == 0:
@@ -244,10 +250,63 @@ def validate_marginals(
             "n_sec_real": n_sec_real,
             "n_sec_pred": n_sec_pred_all,
             "n_sec_accuracy": n_sec_accuracy,
-            "phys_real": phys_real,
-            "phys_generated": phys_gen,
-            "phys_kl": phys_kl,
             "energy_fraction_kl": energy_fraction_kl,
         }
     )
+
+    if target == "physical":
+        phys_real = np.concatenate(all_phys_real, axis=0)  # (M, 2)
+        phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
+
+        if len(phys_real) > 0 and len(phys_gen) > 0:
+            phys_kl = np.array(
+                [
+                    _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
+                    for j in range(2)
+                ]
+            )
+        else:
+            phys_kl = np.full(2, np.nan)
+
+        print(
+            f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
+            f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
+        )
+        print("-" * 68)
+        for j, name in enumerate(_SEC_PHYS_NAMES):
+            r, g = phys_real[:, j], phys_gen[:, j]
+            if len(r) == 0 or len(g) == 0:
+                continue
+            print(
+                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
+                f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
+            )
+
+        result.update(
+            {"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl}
+        )
+    else:
+        type_class_real = np.concatenate(all_type_class_real, axis=0)
+        type_class_gen = np.concatenate(all_type_class_gen, axis=0)
+        n_classes = (
+            sec_decoder.type_dim
+            if target == "onehot"
+            else sec_decoder.cond_enc.pdg_emb.weight.size(0)
+        )
+        type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
+
+        print(
+            f"\n{'sec type class (' + target + ')':<24} "
+            f"n={len(type_class_real)}/{len(type_class_gen)} "
+            f"KL(real||gen)={type_class_kl:.4f}"
+        )
+
+        result.update(
+            {
+                "type_class_real": type_class_real,
+                "type_class_gen": type_class_gen,
+                "type_class_kl": type_class_kl,
+            }
+        )
+
     return result

scripts/dwarf.py

@@ -522,10 +522,21 @@ def warm_cache(
             "--seed", "-s", help="Must match the `giant train` run(s) to warm for"
         ),
     ] = 0,
-    conditioning: Annotated[
+    particle_conditioning: Annotated[
         Conditioning,
         typer.Option(
-            "--conditioning", help="Must match the `giant train` run(s) to warm for"
+            "--particle-conditioning",
+            help="Must match the `giant train` run(s)' conditioning.particle.type "
+            "to warm for",
+        ),
+    ] = Conditioning.physical,
+    material_conditioning: Annotated[
+        Conditioning,
+        typer.Option(
+            "--material-conditioning",
+            help="Must match the `giant train` run(s)' conditioning.material.type "
+            "to warm for — independent of --particle-conditioning "
+            "(docs/v0.3.0-design.md §3.1: the two axes may differ)",
         ),
     ] = Conditioning.physical,
     router: Annotated[
@@ -552,15 +563,17 @@ def warm_cache(
     """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
 
     Warms the vocab maps, event-id split index, and the normalizer entry for
-    the given --val-fraction/--seed/--conditioning, so a later `giant train`
-    run (or a `dwarf hparam-scan` sweep, which shares one such entry across
-    every run) skips straight to training. See giant/data/setup_cache.py.
+    the given --val-fraction/--seed/--particle-conditioning/
+    --material-conditioning, so a later `giant train` run (or a `dwarf
+    hparam-scan` sweep, which shares one such entry across every run) skips
+    straight to training. See giant/data/setup_cache.py.
     """
     run_warm_setup_cache(
         data=str(data),
         val_fraction=val_fraction,
         seed=seed,
-        conditioning=conditioning.value,
+        particle_conditioning=particle_conditioning.value,
+        material_conditioning=material_conditioning.value,
         router_enabled=router,
         router_type=router_type,
         n_experts=n_experts,

scripts/warm_setup_cache.py

@@ -9,6 +9,7 @@ for the sidecar itself.
 
 from pathlib import Path
 
+from giant.constants import K_MAX
 from giant.pipeline import run_setup_stage
 
 
@@ -16,7 +17,8 @@ def run_warm_setup_cache(
     data: str,
     val_fraction: float = 0.1,
     seed: int = 0,
-    conditioning: str = "physical",
+    particle_conditioning: str = "physical",
+    material_conditioning: str = "physical",
     router_enabled: bool = False,
     router_type: str = "energy",
     n_experts: int = 4,
@@ -25,10 +27,12 @@ def run_warm_setup_cache(
 ) -> None:
     """Populate (or refresh) the setup cache sidecar for `data`.
 
-    `val_fraction`/`seed`/`conditioning` select the normalizer cache entry
+    `val_fraction`/`seed`/`particle_conditioning`/`material_conditioning`
+    select the normalizer cache entry
     (`giant.data.setup_cache.normalizer_key`) — pass the same values a later
-    `giant train` invocation will use so it hits this warmed entry.
-    `router_enabled`/`router_type`/`n_experts` only matter for
+    `giant train` invocation will use so it hits this warmed entry. The two
+    conditioning axes are independent (docs/v0.3.0-design.md §3.1) and may
+    differ. `router_enabled`/`router_type`/`n_experts` only matter for
     `router_type == "process"` (warms that `n_experts`'s process map); the
     energy-router quantile summary is always collected regardless, so a
     later `--router-type energy` run never needs to rescan just to seed
@@ -40,16 +44,17 @@ def run_warm_setup_cache(
         "n_experts": n_experts,
     }
     # A minimal v0.3 cfg — only the keys run_setup_stage actually reads
-    # (conditioning.particle.type, stage{1,2}_model.router). This CLI only
-    # ever configures one router (matching today's single --router-type
-    # flag), so it's placed on stage1_model; stage2_model's stays disabled.
+    # (conditioning.{particle,material}.type, stage{1,2}_model.router). This
+    # CLI only ever configures one router (matching today's single
+    # --router-type flag), so it's placed on stage1_model; stage2_model's
+    # stays disabled.
     cfg = {
         "conditioning": {
-            "particle": {"type": conditioning},
-            "material": {"type": conditioning},
+            "particle": {"type": particle_conditioning},
+            "material": {"type": material_conditioning},
         },
         "stage1_model": {"router": router_cfg},
-        "stage2_model": {"router": {"enabled": False}},
+        "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
     }
     run_setup_stage(
         Path(data),

tests/test_cli_predict.py

@@ -1,73 +1,14 @@
 import uuid
 
-import pytest
-import typer
 import yaml
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
-    _check_conditioning_onehot_support,
     _resolve_prediction_output,
     _write_prediction_ref,
 )
 
 
-# ---------------------------------------------------------------------------
-# _check_conditioning_onehot_support
-# ---------------------------------------------------------------------------
-
-
-def _nested_model_cfg(
-    particle_type="physical", material_type="physical", target="physical"
-):
-    return {
-        "conditioning": {
-            "particle": {"type": particle_type, "emb_dim": 8},
-            "material": {"type": material_type, "emb_dim": 8},
-        },
-        "stage2_model": {"particle_type": {"target": target}},
-    }
-
-
-def test_check_conditioning_onehot_support_allows_physical():
-    _check_conditioning_onehot_support(_nested_model_cfg(), "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_rejects_onehot_particle_conditioning():
-    cfg = _nested_model_cfg(particle_type="onehot")
-    with pytest.raises(typer.Exit):
-        _check_conditioning_onehot_support(cfg, "predict")
-
-
-def test_check_conditioning_onehot_support_rejects_onehot_material_conditioning():
-    cfg = _nested_model_cfg(material_type="onehot")
-    with pytest.raises(typer.Exit):
-        _check_conditioning_onehot_support(cfg, "rollout")
-
-
-def test_check_conditioning_onehot_support_allows_onehot_particle_type_target():
-    """stage2_model.particle_type.target="onehot" is implemented (v0.3.0
-    step 6, giant.rollout.decode_secondary_identity) — it's a separate axis
-    from conditioning.particle.type, which this guard doesn't gate at all."""
-    cfg = _nested_model_cfg(target="onehot")
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_allows_embedding_particle_type_target():
-    cfg = _nested_model_cfg(
-        particle_type="embedding", material_type="embedding", target="embedding"
-    )
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_is_noop_for_v02_flat_model_config():
-    """A v0.2 checkpoint's flat model_config has conditioning as a plain
-    string, not a dict — never onehot, so this must be a silent no-op rather
-    than crash on `.get("particle")` against a string."""
-    cfg = {"conditioning": "embedding", "mode": "flow"}
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
 # ---------------------------------------------------------------------------

tests/test_cli_train_overrides.py

@@ -0,0 +1,97 @@
+"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md
+decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must
+independently override each stage's config block, and must take precedence
+over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
+apply the same value to both stages for backward compatibility."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+import giant.cli as cli
+
+runner = CliRunner()
+
+
+def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict:
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["cfg"] = cfg
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args,
+    )
+    assert result.exit_code == 0, result.output
+    return captured["cfg"]
+
+
+def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--mode", "wgan", "--stage1-generator", "flow"],
+    )
+    assert cfg["stage1_model"]["generator"] == "flow"
+    assert cfg["stage2_model"]["generator"] == "wgan"
+
+
+def test_stage2_only_knobs(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        [
+            "--stage2-decoder",
+            "one_shot",
+            "--stage2-k-max",
+            "8",
+            "--stage2-hidden-dim",
+            "32",
+            "--stage2-context-dim",
+            "16",
+            "--stage2-stage1-context",
+            "sampled",
+        ],
+    )
+    assert cfg["stage2_model"]["decoder"] == "one_shot"
+    assert cfg["stage2_model"]["k_max"] == 8
+    assert cfg["stage2_model"]["hidden_dim"] == 32
+    assert cfg["stage2_model"]["context_dim"] == 16
+    assert cfg["stage2_model"]["stage1_context"] == "sampled"
+    # untouched stage1 defaults
+    assert cfg["stage1_model"]["hidden_dim"] == 256
+
+
+def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--hidden-dim", "64", "--stage1-hidden-dim", "128"],
+    )
+    assert cfg["stage1_model"]["hidden_dim"] == 128
+
+
+def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        [
+            "--mode",
+            "wgan",
+            "--n-critic",
+            "5",
+            "--stage1-n-critic",
+            "3",
+            "--stage2-gp-weight",
+            "2.5",
+        ],
+    )
+    assert cfg["stage1_model"]["wgan"]["n_critic"] == 3
+    assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
+    assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
+    assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5

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"
@@ -582,21 +584,19 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
     gconfig.validate_config(cfg)  # must not raise
 
 
-def test_validate_config_mixed_particle_material_conditioning_not_supported():
-    """The data pipeline doesn't support mixed conditioning types yet, even
-    though ConditionEncoder itself already can (docs/v0.3.0-design.md §3.1
-    vs. giant/data/transforms.py's still-single conditioning param)."""
+def test_validate_config_mixed_particle_material_conditioning_is_valid():
+    """docs/v0.3.0-design.md §3.1: the particle and material conditioning
+    axes are configured independently and may mix freely — e.g. material
+    "physical" with particle "embedding" — and the data pipeline
+    (giant/data/transforms.py) now implements that end-to-end, so
+    validate_config must not reject it."""
     cfg = _cfg_with(
         **{
             "conditioning.particle.type": "physical",
             "conditioning.material.type": "embedding",
         }
     )
-    try:
-        gconfig.validate_config(cfg)
-        assert False, "expected ValueError"
-    except ValueError as e:
-        assert "mixed" in str(e) or "conditioning.material.type" in str(e)
+    gconfig.validate_config(cfg)  # must not raise
 
 
 def test_validate_config_pdg_router_incompatible_with_physical_conditioning():
@@ -637,6 +637,49 @@ def test_validate_config_stop_token_not_implemented():
         assert "stop_token" in str(e)
 
 
+def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
+    """docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a
+    rollout-capable checkpoint' — both stages active means giant rollout
+    could load this checkpoint, but 'truth' has no ground truth to draw
+    n_sec from at rollout time."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": True,
+            "stage2_model.active": True,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_sec.mode" in str(e) and "truth" in str(e)
+
+
+def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint():
+    """'truth' is exactly the standalone stage-2 evaluation mode the design
+    doc carves out — stage1_model.active = false must still pass."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": False,
+            "stage2_model.active": True,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_n_sec_truth_allowed_when_stage2_inactive():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": True,
+            "stage2_model.active": False,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 def test_validate_config_ar_default_markov_always_passes():
     """DEFAULT_CONFIG already has decoder='autoregressive',
     history='markov', teacher_forcing='always' — must not raise (v0.3.0
@@ -645,13 +688,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 +726,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:

tests/test_dataset.py

@@ -126,7 +126,8 @@ def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
             target_normalizer=tgt_norm,
             batch_size=4,
             shuffle=False,
-            conditioning="embedding",
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
         )
         return sum(len(batch[0]) for batch in ds)
 

tests/test_dwarf.py

@@ -103,7 +103,7 @@ def test_warm_cache_writes_sidecar(tmp_path):
     assert loaded is not None
     assert loaded.vocab is not None
     assert loaded.event_index is not None
-    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
+    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
 
 
 def test_warm_cache_second_run_hits_cache(tmp_path):
@@ -167,5 +167,5 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
     assert "fitting normalizer (streaming)" in result.output
     loaded = setup_cache.load(data, [data])
     assert loaded is not None
-    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
-    assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers
+    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
+    assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers

tests/test_flow.py

@@ -52,7 +52,7 @@ def test_sample_flow_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
     assert sample.shape == (B, 9)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)
 
 
 def test_ddpm_loss_nonneg():
@@ -69,4 +69,4 @@ def test_sample_ddim_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
     assert sample.shape == (B, 9)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)

tests/test_migration_v02_v03.py

@@ -198,6 +198,56 @@ def test_migrate_legacy_model_config_shape():
     assert migrated["stage2_model"]["decoder"] == "one_shot"
 
 
+def test_migrate_legacy_model_config_nonzero_expert_dims_raises():
+    """docs/v0.3.0-followups.md item 8 regression: a v0.2 checkpoint's
+    model_config carrying a non-default expert_hidden_dim/expert_n_blocks
+    must fail loudly through this path too (§4.2) — not just
+    giant.config.migrate_config's parallel TOML-load path. Silently dropping
+    these keys (build_router's kwarg filtering) would resize the experts
+    instead of refusing."""
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 128,
+        "expert_n_blocks": 0,
+    }
+    try:
+        net._migrate_legacy_model_config(legacy_cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "expert_hidden_dim" in str(e)
+
+
+def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently():
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 0,
+        "expert_n_blocks": 0,
+    }
+    migrated = net._migrate_legacy_model_config(legacy_cfg)
+    assert "expert_hidden_dim" not in migrated["stage1_model"]["router"]
+    assert "expert_n_blocks" not in migrated["stage1_model"]["router"]
+    assert "expert_hidden_dim" not in migrated["stage2_model"]["router"]
+    assert "expert_n_blocks" not in migrated["stage2_model"]["router"]
+
+
+def test_build_models_with_legacy_config_nonzero_expert_dims_raises():
+    """The same check must also fire through the actual caller,
+    build_models, not just the internal helper directly."""
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 128,
+        "expert_n_blocks": 0,
+    }
+    try:
+        net.build_models(legacy_cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "expert_hidden_dim" in str(e)
+
+
 def test_build_models_accepts_new_nested_shape_unchanged():
     """A dict that already has a 'stage1_model' key (the new shape) is
     passed through build_models without going through the legacy migration

tests/test_network.py

@@ -1,13 +1,18 @@
+import copy
+
 import pytest
 import torch
+from giant import config as gconfig
 from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
+    AttentionHistory,
     ConditionEncoder,
     MarkovHistory,
     SinusoidalEmbedding,
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    build_models,
     cat_col_layout,
     stage2_trunk_sec_dim,
     stage2_type_dim,
@@ -323,6 +328,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 +430,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 +590,94 @@ 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
+
+
+# ── build_models: conditioning.share_stages ─────────────────────────────────
+
+
+def _minimal_model_config(share_stages: bool) -> dict:
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["share_stages"] = share_stages
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    return {
+        "pdg_vocab": 3,
+        "mat_vocab": 2,
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def test_build_models_share_stages_true_shares_condition_encoder_instance():
+    built = build_models(_minimal_model_config(share_stages=True))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is stage2.cond_enc
+
+
+def test_build_models_share_stages_false_builds_independent_condition_encoders():
+    built = build_models(_minimal_model_config(share_stages=False))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is not stage2.cond_enc
+
+
+def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists():
+    """The shared encoder's parameters must actually appear in both stages'
+    own `.parameters()` — that's what makes each stage's independent
+    optimizer include (and update) them, which is the actual mechanism behind
+    "shared weights, forced common representation" (docs/v0.3.0-design.md
+    §3.1), not just object identity on `.cond_enc`."""
+    built = build_models(_minimal_model_config(share_stages=True))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+
+    shared_ids = {id(p) for p in stage1.cond_enc.parameters()}
+    assert shared_ids
+    assert shared_ids <= {id(p) for p in stage1.parameters()}
+    assert shared_ids <= {id(p) for p in stage2.parameters()}

tests/test_pipeline.py

@@ -7,6 +7,7 @@ import torch
 
 from giant import config as gconfig
 from giant.data import setup_cache
+from giant.data.transforms import Normalizer
 from giant.pipeline import run_train_job
 
 
@@ -235,6 +236,74 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
     assert "fitting normalizer (streaming)" in joined
 
 
+def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 3 regression: stage2_model.k_max other
+    than the K_MAX module constant's default (15) must not produce a shape
+    mismatch between the data pipeline (loader.py/transforms.py padding) and
+    the model (network.py's trunks, sized from this same config value)."""
+    cfg = _tiny_cfg()
+    cfg["stage2_model"]["k_max"] = 3
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt["model_config"]["stage2_model"]["k_max"] == 3
+
+
+def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type
+    and conditioning.material.type are configured independently and may mix
+    freely (docs/v0.3.0-design.md §3.1) — e.g. particle "embedding" with
+    material "physical" — end-to-end through the real data pipeline, not
+    just accepted by validate_config."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "embedding"
+    cfg["conditioning"]["material"]["type"] = "physical"
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    cond_cfg = ckpt["model_config"]["conditioning"]
+    assert cond_cfg["particle"]["type"] == "embedding"
+    assert cond_cfg["material"]["type"] == "physical"
+
+    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
+    # Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays
+    # unfitted (mean=0/std=1) since "embedding" never computes real values
+    # for it; the material block is fit for real under "physical".
+    from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
+
+    assert cond_norm.mean is not None and cond_norm.std is not None
+    np.testing.assert_allclose(
+        cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
+    )
+    np.testing.assert_allclose(
+        cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
+    )
+    material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
+    assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
+
+
+def test_run_train_job_share_stages_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 5 regression: conditioning.share_stages
+    = true must actually train (not raise NotImplementedError), and the
+    resulting checkpoint's two stages must reload into a single shared
+    ConditionEncoder instance rather than two independent ones."""
+    from giant.model.network import build_models
+
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["share_stages"] = True
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt["model_config"]["conditioning"]["share_stages"] is True
+
+    built = build_models(ckpt["model_config"])
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is stage2.cond_enc
+
+    stage1.load_state_dict(ckpt["model"])
+    stage2.load_state_dict(ckpt["sec_decoder"])
+    for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()):
+        assert torch.equal(p1, p2)
+
+
 def test_run_train_job_matches_uncached_output(tmp_path, data):
     _run(data, tmp_path / "uncached", cache_setup=False)
     _run(data, tmp_path / "cached1", cache_setup=True)

tests/test_rollout.py

@@ -111,7 +111,8 @@ def _run(
         batch_size=128,
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
-        conditioning=conditioning,
+        particle_conditioning=conditioning,
+        material_conditioning=conditioning,
     )
 
 
@@ -172,7 +173,7 @@ def test_seed_frontier_embedding_mode_skips_unresolvable_pdg_lookup():
     frontier construction — mass/charge are simply zero-filled, unused."""
     seeds = _seeds(3)
     seeds["pdg"] = np.full(3, 999999999, dtype=np.int64)
-    fr, _counts = make_seed_frontier(**seeds, conditioning="embedding")
+    fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding")
     np.testing.assert_array_equal(fr["mass"], 0.0)
     np.testing.assert_array_equal(fr["charge"], 0.0)
 
@@ -384,25 +385,43 @@ def _models_v3(
         noise_dim=8,
     )
     particle_type_cfg = {"target": target}
-    common = dict(
-        pdg_vocab=3,
-        mat_vocab=2,
-        particle_cfg=particle_cfg,
-        material_cfg=material_cfg,
-        hidden_dim=32,
-        n_res_blocks=2,
-        generator=generator2,
-        time_dim=16,
-        noise_dim=8,
-        k_max=k_max,
-        particle_type_cfg=particle_type_cfg,
-        build_n_sec_head=stage2_has_n_sec_head,
-    )
+    # Explicit kwargs rather than a shared **common dict: a dict() call whose
+    # values have heterogeneous types (str/int/dict/bool) widens under static
+    # analysis to dict[str, <big union>], which then makes every constructor
+    # keyword not itself part of that union (router, cond_enc, ...) look like
+    # a type mismatch to `ty` even though every actual value passed is fine.
     if decoder == "one_shot":
         sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim)
-        s2 = Stage2OneShot(sec_dim=sec_dim, **common)
+        s2 = Stage2OneShot(
+            pdg_vocab=3,
+            mat_vocab=2,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=32,
+            n_res_blocks=2,
+            generator=generator2,
+            time_dim=16,
+            noise_dim=8,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            build_n_sec_head=stage2_has_n_sec_head,
+            sec_dim=sec_dim,
+        )
     else:
-        s2 = Stage2Autoregressive(**common)
+        s2 = Stage2Autoregressive(
+            pdg_vocab=3,
+            mat_vocab=2,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=32,
+            n_res_blocks=2,
+            generator=generator2,
+            time_dim=16,
+            noise_dim=8,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            build_n_sec_head=stage2_has_n_sec_head,
+        )
     return s1.eval(), s2.eval()
 
 
@@ -440,7 +459,8 @@ def _run_v3(
         batch_size=128,
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
-        conditioning=conditioning,
+        particle_conditioning=conditioning,
+        material_conditioning=conditioning,
         pdg_topn_map=pdg_topn_map,
         other_policy=other_policy,
         seed=seed,
@@ -535,6 +555,88 @@ def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
         _run_v3(s1, s2, pdg_topn_map=None)
 
 
+# --- conditioning.{particle,material}.type = "onehot" (docs/v0.3.0-followups.md
+# item 7) — a separate axis from stage2_model.particle_type.target above: this
+# is what feeds cond_cat's extra top-N columns for ConditionEncoder's own
+# "onehot" mode, not the secondary-species decode. ---------------------------
+
+COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={})
+COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={})
+
+
+def _onehot_conditioning_models():
+    particle_cfg = {"type": "onehot", "emb_dim": len(PDG_MAP), "n_layers": 1}
+    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    )
+    s2 = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3),
+        generator="flow",
+        time_dim=16,
+    )
+    return s1.eval(), s2.eval()
+
+
+def _run_onehot_conditioning(
+    pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP
+):
+    s1, s2 = _onehot_conditioning_models()
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=1.0,
+        max_steps=30,
+        steps=4,
+        batch_size=128,
+        max_tracks_per_event=300,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
+    )
+
+
+def test_rollout_conditioning_onehot_end_to_end(fake_material_props):
+    rec = _run_onehot_conditioning()
+    assert len(rec["event_id"]) > 0
+    assert set(rec["event_id"].tolist()) == set(range(6))
+
+
+def test_rollout_conditioning_onehot_particle_missing_topn_map_raises(
+    fake_material_props,
+):
+    with pytest.raises(RuntimeError, match="pdg_topn_map"):
+        _run_onehot_conditioning(pdg_topn_map=None)
+
+
+def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
+    fake_material_props,
+):
+    with pytest.raises(RuntimeError, match="mat_topn_map"):
+        _run_onehot_conditioning(mat_topn_map=None)
+
+
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 def test_rollout_embedding_target_end_to_end(decoder):
     """particle_type.target="embedding" L1-snaps to the nearest row of the

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()
 
 
@@ -146,7 +150,7 @@ def test_sample_flow_returns_n_sec_for_legacy_stage1():
     )
     cond_cont, cond_cat = _cond(5)
     _, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
-    assert n_sec.shape == (5,)
+    assert n_sec is not None and n_sec.shape == (5,)
 
 
 # ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────
@@ -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,))

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():

tests/test_transforms.py

@@ -426,7 +426,13 @@ def test_build_features_embedding_mode_zero_fills_physical_columns():
     data = _minimal_step_data(3)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding")
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
+    )
 
     assert cond_cont.shape[1] == COND_DIM
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
@@ -438,7 +444,13 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
     data = _minimal_step_data(3)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical")
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
 
     assert cond_cont.shape[1] == COND_DIM
     mass, charge = particle_mass_charge(11)
@@ -461,7 +473,13 @@ def test_build_features_physical_mode_unfilled_material_raises():
     pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0}
 
     with pytest.raises(MaterialPropertiesNotFilledError):
-        build_features(data, pdg_map, mat_map, conditioning="physical")
+        build_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="physical",
+            material_conditioning="physical",
+        )
 
 
 def test_build_cond_features_mass_charge_override(fake_material_props):
@@ -474,7 +492,13 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
     data["charge"] = np.array([2.0, -2.0], dtype=np.float32)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical")
+    cond_cont, _ = build_cond_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
 
     np.testing.assert_allclose(
         cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
@@ -494,7 +518,12 @@ def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
     legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
 
     cond_cont, _ = build_cond_features(
-        data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding"
+        data,
+        pdg_map,
+        mat_map,
+        cond_normalizer=legacy_norm,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
     )
 
     assert cond_cont.shape[-1] == COND_DIM
@@ -521,7 +550,8 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
             pdg_map,
             mat_map,
             cond_normalizer=legacy_norm,
-            conditioning="physical",
+            particle_conditioning="physical",
+            material_conditioning="physical",
         )
 
 
@@ -610,13 +640,23 @@ def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_materi
     }
 
     cond_cont, cond_cat = build_cond_features(
-        data, pdg_map, mat_map, conditioning="physical"
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
     )
     assert cond_cont.shape[-1] == COND_DIM
     np.testing.assert_array_equal(cond_cat, [[0, 0]])  # dummy indices, no raise
 
     with pytest.raises(KeyError):
-        build_cond_features(data, pdg_map, mat_map, conditioning="embedding")
+        build_cond_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
+        )
 
 
 # ── _WelfordAccumulator ──────────────────────────────────────────────────────

tests/test_validate.py

@@ -1,15 +1,18 @@
 import numpy as np
 import torch
 
-from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
-from giant.model.network import Stage1Model, Stage2OneShot
+from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
 from giant.validate import validate_marginals
 
 _PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 _MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+_K_MAX = 5
 
 
-def _tiny_models():
+def _tiny_models(particle_type_cfg: dict | None = None):
+    """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head (decision 1), so
+    n_sec always comes from Stage2OneShot."""
     s1 = Stage1Model(
         pdg_vocab=3,
         mat_vocab=2,
@@ -17,7 +20,13 @@ def _tiny_models():
         material_cfg=_MATERIAL_CFG,
         hidden_dim=16,
         n_res_blocks=1,
-        n_sec_head_k_max=K_MAX,
+    )
+    target = (particle_type_cfg or {}).get("target", "physical")
+    sec_dim = stage2_trunk_sec_dim(
+        particle_type_cfg or {"target": "physical"},
+        "flow",
+        _K_MAX,
+        int(_PARTICLE_CFG["emb_dim"]),
     )
     s2 = Stage2OneShot(
         pdg_vocab=3,
@@ -28,23 +37,29 @@ def _tiny_models():
         n_res_blocks=1,
         generator="flow",
         time_dim=16,
+        k_max=_K_MAX,
+        sec_dim=sec_dim,
+        particle_type_cfg=particle_type_cfg,
     )
+    assert s2.particle_type_cfg.get("target", "physical") == target
     return s1.eval(), s2.eval()
 
 
-def _zero_secondaries_loader(B=4, n_batches=2):
-    """A val_loader whose every batch has n_sec=0 (real side) — matches the
-    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape
-    StreamingStepsDataset yields."""
+def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
+    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
+    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
     batches = []
     for _ in range(n_batches):
         cond_cont = torch.randn(B, COND_DIM)
         cond_cat = torch.zeros(B, 2, dtype=torch.long)
         x1 = torch.randn(B, X_DIM)
-        n_sec = torch.zeros(B, dtype=torch.long)
-        sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM)
+        n_sec = torch.full((B,), n_sec_value, dtype=torch.long)
+        sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx))
+        sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
+        batches.append(
+            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
+        )
     return batches
 
 
@@ -52,22 +67,56 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
     """If n_sec_pred collapses to 0 across the whole validated set (realistic
     during early/unstable training), phys_kl must degrade to NaN instead of
     crashing on the empty-array .min()/.max() reduction inside
-    _histogram_kl -- a regression the old species/bincount code this
-    replaced explicitly guarded against."""
+    _histogram_kl."""
     s1, s2 = _tiny_models()
-    loader = _zero_secondaries_loader()
+    loader = _loader(n_sec_value=0)
 
-    # Force the Stage-1 n_sec head's prediction to 0 for every sample too, so
-    # the generated side's valid-slot mask is also empty (real side is
-    # already all n_sec=0 by construction of the fake loader above).
-    def _fake_sample_flow(model, cond_cont, cond_cat, **kw):
-        B = cond_cont.size(0)
-        return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long)
+    def _fake_resolve_n_sec(
+        stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
+    ):
+        return torch.zeros(cond_cont.size(0), dtype=torch.long)
 
-    monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow)
+    monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)
 
-    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2)
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=2)
 
     assert np.asarray(result["phys_real"]).shape == (0, 2)
     assert np.asarray(result["phys_generated"]).shape == (0, 2)
     assert np.isnan(np.asarray(result["phys_kl"])).all()
+
+
+def test_validate_marginals_physical_target_shapes():
+    s1, s2 = _tiny_models()
+    loader = _loader(n_sec_value=2)
+
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
+
+    assert np.asarray(result["real"]).shape == (4, X_DIM)
+    assert np.asarray(result["generated"]).shape == (4, X_DIM)
+    assert np.asarray(result["kl_divergence"]).shape == (X_DIM,)
+    assert "phys_real" in result and "phys_generated" in result and "phys_kl" in result
+    assert "type_class_real" not in result
+
+
+def test_validate_marginals_onehot_type_class_marginal():
+    particle_type_cfg = {"target": "onehot"}
+    s1, s2 = _tiny_models(particle_type_cfg)
+    loader = _loader(n_sec_value=2, n_classes=s2.type_dim)
+
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
+
+    assert "phys_real" not in result
+    # Real/generated valid-slot counts need not agree (real: ground-truth
+    # n_sec=2 always; generated: the untrained n_sec_head's own prediction).
+    assert np.asarray(result["type_class_real"]).ndim == 1
+    assert np.asarray(result["type_class_gen"]).ndim == 1
+    assert np.asarray(result["type_class_real"]).shape[0] > 0
+
+
+def test_validate_marginals_without_sec_decoder_returns_stage1_only():
+    s1, _ = _tiny_models()
+    loader = _loader(n_sec_value=0)
+
+    result = validate_marginals(s1, loader, n_batches=1, steps=2)
+
+    assert set(result) == {"real", "generated", "kl_divergence"}

tests/test_wgan.py

@@ -120,7 +120,7 @@ def test_sample_wgan_shape():
     cond_cont, cond_cat = _cond(B)
     sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
     assert sample.shape == (B, X_DIM)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)
 
 
 # --- Stage-2 generator/critic ---