Fix/dynamics hooks and order

.claude/skills/nvalchemi-dynamics-api/SKILL.md

@@ -151,9 +151,18 @@ fused = FusedStage(
     exit_status=2,           # auto-set to len(sub_stages) if -1
     compile_step=False,      # enable torch.compile
     compile_kwargs=None,     # kwargs for torch.compile
+    reprime_on_entry=None,   # status codes needing a fresh force compute on entry
 )
 ```
 
+If a sub-stage needs freshly computed forces before it starts integrating
+(e.g. its `AFTER_COMPUTE` hooks differ from the stage the sample came from),
+pass its status code in `reprime_on_entry`. Newly entering samples skip one
+`pre_update`/`post_update` cycle while the shared compute and that stage's
+`AFTER_COMPUTE` hooks refresh their forces, then integrate normally the
+following iteration. This is separate from the initial batch-wide force
+prime and does not affect samples already in that status.
+
 ### With torch.compile
 
 ```python

.claude/skills/nvalchemi-dynamics-hooks/SKILL.md

@@ -77,6 +77,7 @@ class HookContext:
 class DynamicsContext(HookContext):
     step_count: int = 0
     converged_mask: torch.Tensor | None = None
+    active_graph_mask: torch.Tensor | None = None
 ```
 
 Access batch data via `ctx.batch` and dynamics step info via `ctx.step_count`.
@@ -87,9 +88,11 @@ Access batch data via `ctx.batch` and dynamics step info via `ctx.step_count`.
 
 ### Dynamics — `DynamicsStage`
 
-Each `step()` call fires hooks at 9 stages in this order:
+Dynamics exposes 10 lifecycle stages. `ON_ADMISSION` fires once when a
+batch is admitted, while the remaining 9 stages fire within each `step()`:
 
 ```text
+ON_ADMISSION (-1)  ← once before force priming and the first step
 BEFORE_STEP (0)
   BEFORE_PRE_UPDATE (1)  →  pre_update()  →  AFTER_PRE_UPDATE (2)
   BEFORE_COMPUTE (3)     →  compute()      →  AFTER_COMPUTE (4)
@@ -102,12 +105,36 @@ ON_CONVERGE (8)   ← only if convergence detected
 
 | Goal | Stage |
 |------|-------|
+| Validate or allocate for a newly admitted batch | `DynamicsStage.ON_ADMISSION` |
 | Modify forces/energy after model | `DynamicsStage.AFTER_COMPUTE` |
 | Observe final state (logging, snapshots) | `DynamicsStage.AFTER_STEP` |
 | Wrap positions after velocity update | `DynamicsStage.AFTER_POST_UPDATE` |
 | Instrument timing / profiling | `DynamicsStage.BEFORE_STEP` |
 | React to convergence | `DynamicsStage.ON_CONVERGE` |
 
+`ON_ADMISSION` is reset for every new `run()` and for managed membership
+changes such as refill or pipeline communication. In `FusedStage`, it runs
+outside compiled `_step_impl`, making it suitable for shape-dependent allocation
+and Python setup that per-step hooks cannot safely perform under `fullgraph=True`.
+
+In `FusedStage`, fused-level hooks wrap sub-stage hooks at every shared boundary:
+fused `BEFORE_*` hooks run before the corresponding sub-stage loop, and fused
+`AFTER_*` hooks run after it. This includes the pre-update and post-update
+boundaries. Fused hooks receive the overall active mask. The mask passed to a
+sub-stage hook depends on the boundary:
+
+- `BEFORE_STEP`, `BEFORE_COMPUTE`, `AFTER_COMPUTE`, and `AFTER_STEP` receive
+  the sub-stage's status mask captured at the start of the fused step. This
+  includes graphs with `reprime_pending` so compute-boundary hooks can prepare
+  and observe their refreshed model outputs.
+- `BEFORE_PRE_UPDATE`, `AFTER_PRE_UPDATE`, `BEFORE_POST_UPDATE`, and
+  `AFTER_POST_UPDATE` receive the narrower update-eligibility mask: the
+  sub-stage status mask with `reprime_pending` graphs excluded. The same mask
+  brackets both integrator updates, even though force computation clears
+  `reprime_pending` before the post-update boundary.
+- `ON_CONVERGE` receives that sub-stage's convergence mask and remains
+  sub-stage-only because convergence is evaluated independently per sub-stage.
+
 ---
 
 ## Registering hooks

.claude/skills/nvalchemi-dynamics-implementation/SKILL.md

@@ -24,9 +24,11 @@ from nvalchemi.data import Batch
 
 ## Step execution flow
 
-Each call to `step(batch)` executes:
+The first `step(batch)` after admission dispatches `ON_ADMISSION`, while subsequent
+steps skip it until admission is explicitly reset. The per-step sequence is:
 
 ```text
+0. ON_ADMISSION hooks (once after reset, before the compiled step)
 1. BEFORE_STEP hooks
 2. BEFORE_PRE_UPDATE hooks  →  pre_update(batch)  →  AFTER_PRE_UPDATE hooks
 3. BEFORE_COMPUTE hooks     →  compute(batch)      →  AFTER_COMPUTE hooks

CHANGELOG.md

@@ -7,6 +7,31 @@
 - Add support for PEFT fine-tuning within `FineTuningStrategy`, including
   LoRA workflows with `LoRAConfig`, `load_peft_checkpoint_into_model`,
   and base-model fingerprint checks for PEFT checkpoint loading.
+- `FusedStage(reprime_on_entry=...)` — status codes whose newly entering
+  graphs skip one integrator update so the shared compute and target-stage
+  `AFTER_COMPUTE` hooks can refresh forces under the new stage's context
+  before it advances them.
+
+### Fixed
+
+- **Dynamics hook lifecycle** — add `ON_ADMISSION` for one-time batch setup
+  before force priming and outside compiled fused steps. Fused stages now fire
+  shared step and compute hooks at both fused and sub-stage levels with
+  consistent nesting.
+- **`FusedStage` force priming** — a dynamics instance's own adaptive
+  optimizer state (e.g. FIRE's per-graph `dt`/`alpha`/step counters in
+  `self._state`) is now preserved across masked `pre_update`/`post_update` calls.
+
+### Deprecated
+
+- `FusedStage.register_fused_hook()`. Use the inherited `register_hook()`
+  method instead; hooks on a `FusedStage` already observe the complete fused
+  batch.
+
+## 0.2.0 — 2026-08-07
+
+### Added
+
 - Domain decomposition for distributed inference and dynamics: a spatial halo
   strategy and a graph-parallel strategy, both driven by a declarative
   `MLIPSpec` a model wrapper publishes as `distribution_spec`. Ewald, PME,

docs/index.md

@@ -48,7 +48,7 @@ batching and size-aware sampling handle load balancing automatically.
 - **Bring your own model** --- wrap MACE, AIMNet2, or any PyTorch MLIP in a
   few lines with a standardized `ModelConfig` interface.
 - **Compose, don't configure** --- fuse stages on one GPU with `+`, distribute
-  across GPUs with `|`, and inject behavior at nine hook points per step.
+  across GPUs with `|`, and inject behavior at admission and per-step hook points.
 - **GPU-native data** --- `AtomicData` and `Batch` are Pydantic-validated,
   `jaxtyping`-annotated graph structures that live on-device.
 - **Inflight batching** --- converged samples are replaced mid-run so the GPU

docs/modules/dynamics/buffers_and_data_flow.rst

@@ -351,7 +351,7 @@ iterations, and returned.
        edge [fontsize=10]
 
        init [label="1. sampler.build_initial_batch()\nstatus=0, fmax=inf" fillcolor="#1a1a1a"]
-       step [label="2. compute() →\nper-sub-stage masked_update\nbased on batch.status"]
+       step [label="2. masked pre_updates → shared compute() →\nmasked post_updates based on batch.status"]
        conv [label="3. ConvergenceHook\nupdates batch.status\n(0 → 1 → 2 …)"]
        refill [label="4. _refill_check()\nevery refill_frequency steps" fillcolor="#4a3315"]
        refill_detail [label="identify graduated (status ≥ exit_status)\nwrite to sinks · extract remaining\nrequest replacements · rebuild tensors" shape=plaintext fillcolor=none style=""]

docs/modules/dynamics/fused_stage.rst

@@ -70,23 +70,24 @@ from 0 when using ``+``). Every sample in the batch carries a
            fontsize=12
 
            batch   [label="Batch (8 samples)\nstatus: [0, 0, 0, 1, 1, 0, 1, 0]" fillcolor="#4a3315"]
-           compute [label="1. compute()\nsingle forward pass for ALL 8 samples"]
-           mask0   [label="2. sub_stage[0].masked_update\nstatus == 0  →  samples 0, 1, 2, 5, 7"]
-           mask1   [label="3. sub_stage[1].masked_update\nstatus == 1  →  samples 3, 4, 6"]
+           pre0    [label="1a. sub_stage[0]._masked_pre_update\nstatus == 0  →  samples 0, 1, 2, 5, 7"]
+           pre1    [label="1b. sub_stage[1]._masked_pre_update\nstatus == 1  →  samples 3, 4, 6"]
+           compute [label="2. compute()\nsingle shared forward pass for ALL 8 samples"]
+           post0   [label="3a. sub_stage[0]._masked_post_update\nstatus == 0  →  samples 0, 1, 2, 5, 7"]
+           post1   [label="3b. sub_stage[1]._masked_post_update\nstatus == 1  →  samples 3, 4, 6"]
            conv    [label="4. convergence check\nper sub-stage"]
            migrate [label="sample 2 converges → status[2] = 1\n(migrated!)" shape=plaintext fillcolor=none style=""]
 
-           batch -> compute [style=bold]
-           compute -> mask0 [style=bold]
-           mask0 -> mask1 [style=bold]
-           mask1 -> conv [style=bold]
+           batch -> pre0 [style=bold]
+           pre0 -> pre1 -> compute -> post0 -> post1 -> conv [style=bold]
            conv -> migrate [style=dashed color="#999999"]
        }
    }
 
 The key insight is that **only one forward pass happens** regardless of
-how many sub-stages exist. The expensive model evaluation is amortized
-across all stages.
+how many sub-stages exist. Each sub-stage applies its masked ``pre_update()``
+before that shared evaluation and its masked ``post_update()`` afterward. The
+expensive model evaluation is amortized across all stages.
 
 
 Convergence-driven stage migration
@@ -136,10 +137,11 @@ This is separate from the initial batch-wide force prime.
 Running a ``FusedStage``
 -------------------------
 
-``FusedStage.run()`` loops until **all** samples reach ``exit_status``.
-Unlike ``BaseDynamics.run()``, the ``n_steps`` attribute (inherited
-from ``BaseDynamics``) and any ``n_steps`` argument to ``run()`` are
-**unused** — termination is purely convergence-driven.
+``FusedStage.run()`` loops until all samples reach ``exit_status``, the sampler
+is exhausted, or the maximum ``n_steps`` is reached. An ``n_steps`` argument
+overrides ``FusedStage.n_steps``. When both are ``None``, termination is
+convergence- or sampler-driven. A sub-stage's own ``n_steps`` limits how many
+steps each system remains in that sub-stage before moving to the next one.
 
 **Mode 1: external batch (the common case)**
 
@@ -207,6 +209,12 @@ When ``compile_step=True``, the internal ``_step_impl`` method is
 wrapped with ``torch.compile``. This can significantly improve
 throughput by fusing GPU kernels across the entire fused step.
 
+Before entering ``_step_impl``, :class:`FusedStage` initializes bookkeeping and
+sub-stage state, then dispatches ``ON_ADMISSION`` hooks at the fused and sub-stage
+levels. Admission runs once for a new run or managed batch replacement and stays
+outside the compiled graph. Use it for validation, shape-dependent allocation,
+or other Python setup that is not safe in a per-step compiled hook.
+
 
 Combining with hooks and sinks
 -------------------------------