code/docs/ARCHITECTURE_V2_PROPOSAL.md · viai957/align at main

align / code /docs /ARCHITECTURE_V2_PROPOSAL.md
Add files using upload-large-folder tool
d909bd8 verified 16 days ago
82.6 kB
# align architecture v2: proposal

Status: proposal only. No code was moved or edited to write it. Date: 2026-10-14. Scope: `/apps/development/rlhf/align`
at HEAD `1c483cc`, compared with `reference-repos/AgileRL` at `4f07424` (Apache-2.0).

How to read the evidence labels used below.

- **[read]** read in source; the path and line are given and can be re-checked with `sed -n`.
- **[calc]** arithmetic from numbers that are themselves labelled. The arithmetic is shown or reproducible in five lines.
- **[cpu-probe]** a small CPU run done for this document (the probe is described where it is used; nothing was left in the repo).
- **[assumed]** a number taken from a datasheet or a typical value, not measured on this box. Every one is listed in 5.1.
- **[paper]** arXiv id plus section or equation, checked against the text fetched into the hyperresearch vault.

Independent verification pass (2026-10-14, builder A2): the freeze closure, the layer assignment, the shim mechanics, the
source-pinning tests and the Section 5 arithmetic were re-derived from the repo; corrections are folded into the text below and
marked **[A2]**.

Nothing here was measured on a GPU (the P2 chain owns all four). Section 5 is therefore a model with stated assumptions and a
measurement plan, not a benchmark. The CPU test suite was not run for this document either; the three docs that quote a test
count disagree (1089 in DEPENDENCY_POLICY, 1107 in REFERENCE_GAP_ANALYSIS, 1608 in README), so step M0 records the real baseline
first.
---
## 0. Summary of decisions
1. **Adopt from AgileRL: interfaces and contracts, not structure.** Seven ideas are worth taking (1.3): a typed, validated experience
   container; algorithm variants as a thin switch over one implementation (CISPO and GSPO are `GRPO` with a fixed `loss_type`);
   importance-sampling level and advantage granularity as two independent axes; a text-env protocol separated from a token-level
   rollout harness; an async collector; a documented-constructor AST test; a shard-runtime seam for dense versus FSDP2.
2. **Do not adopt its structure or its dependencies.** The 6,398-line `agilerl/algorithms/core/base.py` makes the algorithm own
   generation, the optimizer step, vLLM colocation, FSDP and LoRA plumbing. align's `Trainer` / `TrainModule` split is already
   better. accelerate, transformers, peft, vLLM, liger-kernel, pydantic, hydra, wandb, dill-pickled checkpoints and the
   evolutionary HPO stack are all out (1.4).
3. **The one structural gap in align is the environment.** `GrpoBatchSource` is hard-wired to `{problem, answer}` rows plus an
   outcome verifier; the multi-turn Lean path is a subclass that re-implements `bind`. A single `Env` protocol removes the fork.
4. **Layering is currently violated in 5 places by lazy imports** (2.1; **[A2]** one is an artefact of putting `eval` at L5, two are
   missed ones). They are cheap to fix and a 40-line AST test keeps them fixed.
5. **Migration is facade-first.** New packages (`algorithms/`, `envs/`, `recipes/`, `optim/`) are added as thin layers over the old
   modules. Old modules move only after the P2 chain exits, with identity-preserving `sys.modules` shims (verified in 4.2).
   **[A2]** 105 of the 129 modules in `align/` are in the import closure of the running P2 scripts and of the GPU-validation steps
   the chain scripts launch; they cannot be touched for now (4.1).
6. **Kernels: only one custom Triton kernel pays a lot, and it is not one of the three the task listed first.**
   - Paged-attention decode (`TritonPagedAttention`, an interface with no kernel today): about 2x to 4x on decode steps at batch
     6 to 64 [calc, 5.4-C; **[A2]** rows that exceed the KV capacity of the card were removed]. This is the big one, because rollout generation dominates an RL iteration (5.2).
   - A fused head + policy-gradient loss pays about 4x on the head (about 13% of an RL *training* step) but the gain comes from
     precision path and GEMM count (4 fp32 GEMMs to 2 bf16 GEMMs), **not** from fusing elementwise ops. The elementwise part is
     about 4% of today's head time and 10% to 16% of the fused head's time **[A2]**. Do it as a chunked `autograd.Function` first; Triton is optional there (5.4-A, 5.4-B).
   - Fusing only the per-token CISPO/GRPO loss arithmetic is worth nothing (64 KB tensors).
---

## 1. AgileRL: what to adopt and what not to

### 1.1 What AgileRL is

An evolutionary-HPO RL library (populations, mutation, tournament selection) that grew an LLM fine-tuning half. Size by directory
[read, `find ... | xargs cat | wc -l`]:

| dir | lines | relevance to align |
|---|---|---|
| `algorithms/` | 27,282 | `core/base.py` 6,398; `grpo.py` 2,471; `ppo_llm.py` 1,415; `reinforce_llm.py` 1,125; `dpo.py` 692; `sft.py` 410; `cispo.py` 26; `gspo.py` 23. The rest is classic RL (DQN, TD3, MADDPG, bandits) |
| `utils/` | 17,179 | `llm_utils.py` 3,035, `algo_utils.py` 2,771, probe envs, mutation utils |
| `modules/` + `networks/` + `wrappers/` + `vector/` + `hpo/` + `population.py` | 15,612 | evolvable networks and populations; none of it applies |
| `llm_envs/` + `training/` + `strategies/` + `builders/` + `models/` | 11,496 | the environment and trainer side; partly relevant |
| `distributed/` | 3,112 | FSDP2 runtime, process helpers |
| `algorithms/core/llm_ops/` | fused ops | `fused_logprobs.py`, `fused_loss.py` (Liger), `fused_lora.py`, `moe_lora.py` |

Roughly 40% of the code is classic-RL or evolution machinery that an LLM post-training library does not need. That is the
first reason not to copy its layout wholesale.

### 1.2 How its LLM algorithms are built

**Class tree** [read]. `EvolvableAlgorithm` (`algorithms/core/base.py:436`, metaclass `RegistryMeta` at `:312` runs
`_registry_init` after every constructor) holds mutation registry, metrics, accelerator and device. `LLMAlgorithm`
(`base.py:2687`, about 3,700 lines of the file) adds: actor/reference/critic adapters, vLLM colocation and weight sync
(`_sync_actor_to_vllm` `:5466`, `_generate_with_vllm_colocate` `:5484`), FSDP shard runtime, fused logprob forward
(`_fused_forward` `:4844`, `_fused_packed_forward` `:4915`), LoRA attach (`_attach_lora_adapters` `:4341`), backward and grad clip
(`_backward_pass` `:5283`), checkpointing (`save_checkpoint` `:3157`).

**Algorithms**. `GRPO(LLMAlgorithm)` (`grpo.py:170`; ctor `:398` with about 50 parameters; `get_action` `:541` generates; `learn`
`:634` runs the epoch/minibatch loop, loss, backward and step). `CISPO(GRPO)` (`cispo.py:15`) and `GSPO(GRPO)` (`gspo.py:15`) are
about 25 lines each: they call `super().__init__(..., loss_type="cispo")`. The objective is chosen in `_setup_objective`
(`grpo.py:1071`): `loss_type in {grpo, gspo, cispo}`, `importance_sampling_level in {token, turn, trajectory}` (gspo forces
`trajectory`), `advantage_granularity` separately, `loss_norm in {micro_batch, accumulation_window}`. `DPO` (`dpo.py:44`) and `SFT`
(`sft.py:42`) follow the same `get_action` / `learn(experiences)` shape with a typed `experiences` argument.

**ILQL and BC_LM are a second generation.** `ilql.py:42` is `class ILQL(nn.Module)` of 2,256 lines, a port of an older
offline-RL codebase with its own `Language_Environment` / `RL_Dataset` types, and `import wandb` at module top (`ilql.py:15`). It
does not use `LLMAlgorithm`. So AgileRL itself carries two incompatible LLM-algorithm architectures. That is the cost of
growing a base class instead of defining a contract first.
**Experience container** [read, `components/llm_rollout_data.py`]. `Trajectory` (frozen pydantic model: `token_ids` `(1,T)`,
`action_masks`/`turn_ids` `(1,T-1)`, `rewards` per turn, optional `sampling_logps`, `pixel_values`) with a `model_validator` that
fails when per-token fields are not `T-1` long. `RolloutGroup` requires exactly `group_size` trajectories. They are packed into
`LLMExperienceBatch`. The label-aligned `T-1` convention is the same one align states in `rl/types.py`.
**Environments** [read]. `protocols.py:602` `TextEnvProtocol` (`reset(...)`, `step(action_text)`, `system_prompt`) and `:624`
`EnvClientProtocol` (Gym five-tuple, `dataset_size`, `tools`, `rubric_components`, `eval_mode()`). `RolloutHarness`
(`llm_envs/harness.py:44`, 857 lines) does tokenisation, chat-template boundaries and the turn loop over a client.
`RolloutCollector` (`collector.py:28`) runs many harnesses; `TaskAssigner` (`task_assigner.py:27`) assigns dataset rows to slots;
`AsyncBatchCollector` (`async_collector.py:24`) overlaps collection with learning; `rubrics.py` composes reward components.

**Training loop and config** [read]. `training/llm/rollout.py:100` `train_llm_rollout(pop, max_turns, env_factory, ...)`. Config is
YAML under `configs/training/llm_finetuning/` (cispo, dpo, grpo, gspo, ppo_llm, reinforce_llm, sft) parsed into pydantic
`*Spec` models (`models/`, `builders/llm.py:46` `LLMBuilder`).

**Checkpointing** [read, `base.py:3157-3250`]. A directory with PEFT adapter dirs plus `attributes.pt`, a pickle written with
`dill` (`base.py:33`, `:1366`, `:1381`). FSDP2 state is gathered to full tensors before saving, so checkpoints are rank-count
independent. Only rank 0 writes.

**Distributed** [read]. `accelerate` for process setup, plus a native FSDP2 layer: `distributed/runtime.py` defines a shard
runtime (`prepare_actor`, `export_model_state`, `export_optimizer_state`) with dense and FSDP2 implementations.

**Tests and docs** [read]. `tests/` mirrors the package (`test_algorithms/test_llms`, `test_llm_ops`, `test_llm_envs`, ...).
`test_algorithms/test_param_docstrings.py` is an AST test: every constructor parameter of an audited class has a `:param:` entry
and no stale entries remain. Sphinx docs under `docs/llm_finetuning/` (batch sizing, fused logprobs, checkpoints, environments).

**Convergent design worth noting.** AgileRL's `fused_logprobs.py` (`fused_linear_logprobs_chunked` `:135`,
`FusedLinearLogProbsFunction` `:199`) is the same chunked head + log-softmax + gradient-checkpoint idea as align's
`kernels/fused_ce.py` and `rl/cispo.py::fp32_head_logprobs`. AgileRL's own GRPO docstring says its Liger GRPO kernel "shows no
speedup over AgileRL's already memory-bounded standard path" (`grpo.py`, `use_liger_loss` parameter doc). That matches the
roofline result in 5.4-A: chunking removes the memory problem, and a fused elementwise kernel adds almost nothing on top.

### 1.3 Ideas to adopt

| # | Idea | AgileRL reference | align today | Concrete action |
|---|---|---|---|---|
| A1 | Typed, validated experience batch with a shape-contract validator | `components/llm_rollout_data.py` `Trajectory`, `RolloutGroup` | `rl/types.py` `RolloutBatch.__post_init__` (`:51`) already validates; `TokenAdvRolloutBatch` lives in `rl/grpo_loss.py:301`; `DpoBatch` in `train/dpo_stage.py:818`; SFT batch is an untyped dict from `collate_sft` | One module `align/batch.py` holding `SFTBatch`, `PreferenceBatch`, `RolloutBatch`, `TokenAdvRolloutBatch`, `PPORolloutBatch`, each with `validate()`. Dataclasses, not pydantic |
| A2 | Variants as a thin switch, not copies | `cispo.py` / `gspo.py` (25 lines each) | `GRPOLossType.{dapo,cispo}` in `rl/grpo_loss.py:41` already does it for the loss; GSPO absent | `algorithms/grpo/` owns one implementation; `algorithms/cispo.py` and `dapo.py` are presets (2.3). Add the trajectory-level ratio as a third option only when needed (GSPO, arXiv:2507.18071 Sec. 4.1 Eq. 5-7: `s_i = exp(mean_t log(pi/pi_old))`) |
| A3 | Two independent axes: importance-sampling level and advantage granularity | `grpo.py:1071` `_setup_objective`, `:1493` `_resolve_advantage_granularity` | `rl/multiturn_source.py` computes per-turn advantages, `expand_advantages` (`grpo_loss.py:234`) broadcasts; ratio level is token only | Name both axes in `GRPOConfig` (`ratio_level`, `advantage_level`) and refuse the mismatched pair that AgileRL only warns about (`grpo.py` warns on `turn` advantage with `trajectory` ratio) |
| A4 | Text-env protocol separate from the token-level harness | `protocols.py:602`, `llm_envs/harness.py:44` | `rl/envs/lean_multiturn.py` fuses both; `LeanMultiTurnEnv.run_group(specs, generate)` already takes the generator as a callable, which is the right seam | `align/envs/` (2.4): `Env` protocol over text, one `Harness` for tokenisation and per-turn masks, `MathEnv`, `LeanEnv`, `CodeEnv` |
| A5 | Async collection overlapping learning | `llm_envs/async_collector.py:24` | rollout is synchronous (`rl/batch_source.py` `_produce_step`); REFERENCE_GAP_ANALYSIS 1.3 already lists this | Defer. Design the `RolloutBackend.submit/collect` protocol (already present in `rollout/protocol.py:95`) so a one-step-stale pipeline is a source-level change, not an engine change |
| A6 | Documented-constructor AST test | `tests/test_algorithms/test_param_docstrings.py` | align documents fields as `#:` comments on dataclasses (`config.py`) but nothing checks it | `tests/contract/test_config_docs.py`: every field of every registered algorithm Config has a doc comment. Stdlib `ast`, runs in the firewall lane |
| A7 | Shard-runtime seam (dense vs FSDP2 vs TP) | `distributed/runtime.py` | `TrainModule.full_model_state_dict`, `optim_state_dict`, `load_*` already are this seam (`train/train_module.py`), plus `model/tp.py` for inference TP | Name the seam `ShardRuntime` in `align/protocols.py` and make `TrainModule` and `TPPagedEngine` implement it; no code change to the Trainer |

Also adopted as a habit rather than code: AgileRL keeps fused ops in their own package with their own tests
(`algorithms/core/llm_ops/`, `tests/test_algorithms/test_llm_ops/`). align already does this in `kernels/`.

### 1.4 Ideas not to adopt

| Idea | AgileRL reference | Why not |
|---|---|---|
| A 6,398-line base class that owns generation, vLLM colocation, FSDP, LoRA and the optimizer step | `algorithms/core/base.py` | The algorithm cannot be tested without those. align's `TrainModule` owns model+loss+optimizer, `Trainer` owns loop and state, the batch source owns rollouts. Each is testable on CPU alone (1,000+ CPU tests prove it) |
| `get_action` generates inside the algorithm | `grpo.py:541` | Couples the algorithm to HF `generate` or colocated vLLM. align's `RolloutBackend` (`rollout/protocol.py:95`) keeps the engine replaceable (native, paged, spool/vLLM) |
| Replacing `lm_head` with a pass-through `HiddenStateHead` (documented as `nn.Identity`) by monkeypatch to get hidden states | `base.py:6314` `_patch_lm_head_to_identity` | Mutates the module tree under a context manager. align passes `hidden_consumer(hidden, lm_head_weight)` into `forward` (`model/qwen3.py:467`, `mistral.py:597`): explicit, and it fails loudly on a model without the kwarg (`rl/cispo.py` `fp32_head_forward`) |
| accelerate, transformers, peft, bitsandbytes, liger-kernel, vLLM as hard or core dependencies | `pyproject.toml` `dependencies`, `cpu-llm` extra | Violates `docs/DEPENDENCY_POLICY.md` and `tests/test_dependency_firewall.py`. align already has the replacements (own LoRA, own 8-bit AdamW, own tokenizer, own paged engine) |
| pydantic spec models and hydra/omegaconf | `models/`, `builders/` | `align/config.py` is dataclass + `tomllib` with struct mode and typed dotted overrides, no dependency. Adopt only the idea that the algorithm's config is its own dataclass |
| dill-pickled `attributes.pt` as the checkpoint container | `base.py:33`, `:1366` | Pickle loads arbitrary code. align reads `resume.pt` with `weights_only=True` and has fsync, partial-staging and completeness checks (`train/trainer.py`) AgileRL lacks |
| Evolutionary HPO: population, mutation, tournament, `EvolvableModule` networks | `hpo/`, `population.py`, `modules/`, `networks/` | Irrelevant to single-run post-training, and it shapes the base class (`MutationRegistry`, `_registry_init` metaclass hook) |
| The older `nn.Module` ILQL/BC_LM generation | `algorithms/ilql.py`, `bc_lm.py` | Shows the failure mode of not fixing the contract first (1.2) |
| Arena/cloud manifests | `models/manifest.py`, `agilerl-arena` | Not a use case here |
| Warn-and-continue on conflicting options | many `warnings.warn` in `_setup_objective` | align's convention is a launch refusal (`assert_a3_conformant`, `assert_p4_cispo_conformant`). Keep refusing |

### 1.5 Where align is already ahead (keep as is)

- The Trainer / TrainModule split with `extra_state()` hooks and a documented resume flow (`docs/ARCHITECTURE.md`, "Trainer /
  TrainModule split").
- The control plane: coordinated stop, MAX-reduce decisions, exit codes 80 to 83, `HALT.json`, supervisor
  (`align/exit_codes.py`, `utils/supervisor.py`).
- Host safety: preflight, heavy-job flock, D-state gate (`utils/guard.py`, `utils/host_state_gate.py`).
- The dependency firewall as an executable test.
- Launch-refusal gates per stage with a dated rationale in each docstring.
- Bit-invariance as a first-class engine mode (`TorchPagedAttention(mode="row")`).

None of this exists in AgileRL, and none of it should be traded for its structure.

Licence note: AgileRL is Apache-2.0. Adopting ideas needs no attribution; copying code needs the header and a NOTICE entry. This
proposal assumes ideas only.

---

## 2. Target package layout

### 2.1 Layer rules, and where the current code breaks them

Imports flow downward only. A layer may import its own layer only through an explicit allow-list in the test.

```
L0  root modules: config, exceptions, exit_codes, registry(banner), protocols (new), batch (new)
    runtime/                     (was utils/)
L1  kernels, tokenizer, model, optim (new), functional (new: pure tensor/scalar math, no align imports)
L2  verify, data, rollout, inference, envs (new)
L2b eval             (imports only L1/L2: data, verify, inference, model, tokenizer; measured, see below)
L3  train            (the engine: Trainer, TrainModule, checkpoint formats, FSDP, stage gates and smoke)
L4  algorithms (new) (losses, batch contracts, modules, presets; subclass TrainModule)
L5  recipes (new), cli (new)
```

**[A2] Why `eval` is L2b and not L5.** The first draft put `eval` at L5 and then listed `train -> eval` as an upward import. An AST
scan of `align/eval/*.py` shows its only `align` imports are `config`, `data.schema`, `data.tokenization`, `inference.engine`,
`model.*`, `tokenizer`, `verify` and itself, so `eval` sits below `train`. With `eval` at L5 the proposed fixes created two new
violations (moving `train/stage_gates.py` to `recipes/` makes `dpo_stage`, now L4, import L5; moving `ref_cache` to
`algorithms/dpo/` leaves `ref_cache -> eval.metrics.wilson_interval`, L4 to L5) and moved a 955-line, 12-import-site, frozen file
for no gain. With `eval` at L2b, `stage_gates.py`, `stage_smoke.py` and `ref_cache.py` need no layering fix. The `align/__init__.py`
lazy public API is exempt from the layer rule (it names L2 modules from L0 through `__getattr__`).

Edges measured with an AST scan of every `import` including lazy ones [read, script run for this document]:

| Violation | Evidence | Fix |
|---|---|---|
| ~~`train` imports `eval`~~ **[A2] not a violation** once `eval` is L2b (see above) | `train/stage_smoke.py:369,370,453,463`, `train/ref_cache.py:594`, `train/stage_gates.py:351,685` | none: `stage_gates`, `stage_smoke` stay in `train/`; `ref_cache` moves to `algorithms/dpo/` and may import `eval` |
| `train` imports `ppo` (upward) | `train/qwen3_stage.py:171,189,209,224,258` (`ppo.lora`) and `:333` (`ppo.trainer.AdapterProbeView`) **[A2: the first draft listed only 171,189,209 and missed the `AdapterProbeView` edge]** | LoRA is a model concern: `ppo/lora.py` to `model/lora.py`; `AdapterProbeView` is a parameter view for `train.probes.TrunkMovementProbe`, so it moves next to it, `ppo/trainer.py::AdapterProbeView` to `train/probes.py` |
| **[A2] new:** `rl/envs/lean_multiturn.py` (becomes `envs/`, L2) imports `rl.advantages` and `rl.rewards` (would be `algorithms/`, L4) | `lean_multiturn.py:100,101` (`compute_group_advantages`, `soft_overlong_punishment`) | `rl/advantages.py`, `rl/rewards.py`, `rl/ops.py`, `ppo/gae.py` have no `align` imports (checked); they become `functional/{advantages,rewards,ops,gae}.py` at L1 instead of `algorithms/common/` |
| `data` imports `eval` while `eval` imports `data` (cycle) | `data/outcome_agreement.py:275,306` imports `eval.verifier`; `eval` has 3 edges into `data` | `Outcome`/`OutcomeVerifier` to `verify/outcome.py`; `outcome_agreement.py` to `eval/` |
| `verify` imports `eval` while `eval` imports `verify` (cycle) | `verify/lean/reward.py:116` imports `eval.verifier.Outcome` | same move |
| The library loads code from `scripts/` and `extras/` by file path | `rl/envs/lean_multiturn.py:108-123` (`_load_by_path("...", "extras/vllm_rollout/lean_prompts.py")`, `"scripts/p2/p2_common.py"`) | Vendor the prompt module as `envs/lean/prompts.py` with a sha256 identity test against the `extras/` copy; extract the one function used from `p2_common` (`displayed_statement_matches`) |
| Hard-coded absolute paths in library code | `config.py:40,335,339`, `kernels/gate.py:36`, `kernels/_toolchain.py:42`, `utils/guard.py:41,42,43` (`PROJECT_ROOT = "/apps/development/rlhf/align"`); `eval/golden.py:84,123,128,129,374`; `verify/lean/toolchain.py:39`; `rollout/client.py:135`; `data/hub.py:54`, `data/sources.py:37`, `data/ingest_parquet.py:36`; `train/dpo_stage.py:142,149` (**[A2]** the last three groups were missing; `config.py:335,339` too) | M8; blockers for open-sourcing. **Editing `eval/golden.py` changes `eval_source_hash` (4.5)** |

The enforcing test (`tests/test_layering.py`, about 40 lines) builds the same graph, maps each module to a layer from a table in
the test, and fails on an upward edge or on a cycle between two same-layer packages. **[A2]** It scans only real modules: files
listed in the shim manifest (4.2) are skipped, and the legacy packages `ppo`, `rl`, `rl.envs`, `utils` are treated as aliases of their
targets, otherwise every shim would be reported as an upward edge. It starts in report-only mode (M0) so the
current violations are recorded, then becomes strict as each is fixed.

### 2.2 Tree

```
align/
  __init__.py                 lazy public API (existing pattern; new names added in M7)
  config.py exceptions.py exit_codes.py registry.py            unchanged location (L0)
  protocols.py                NEW  Algorithm, Env, RolloutBackend, Verifier, ShardRuntime, WeightFormat, LossAdapter
  batch.py                    NEW  SFTBatch, PreferenceBatch, RolloutBatch, TokenAdvRolloutBatch, PPORolloutBatch
                              ([A2] both files import torch only under TYPE_CHECKING or inside functions; neither is imported by
                              `align/__init__`, and `test_public_api`'s heavy-module list gains `align.batch`, `align.protocols`)
  runtime/                    dist.py guard.py host_state_gate.py supervisor.py tracking.py   (from utils/)
                              __init__.py stays empty: `import align.runtime` must not import torch (tests/test_public_api.py)
  functional/                 NEW (L1)  advantages.py rewards.py ops.py gae.py kl.py   pure math, no align imports
  kernels/                    unchanged + fused_policy_loss.py (NEW), paged_attention.py (NEW, Triton)
  tokenizer/                  unchanged
  model/                      unchanged + lora.py (from ppo/lora.py), heads.py (from ppo/value_head.py),
                              build.py (arch registry: mistral | qwen3 | olmoe)
  optim/                      adamw8bit.py quant8.py offload_adamw.py factory.py (make_optimizer)
  verify/                     unchanged + outcome.py (from eval/verifier.py), code/ (NEW)
  data/                       unchanged (minus outcome_agreement.py)
  rollout/                    unchanged + backends.py (NEW registry: native | paged | tp | spool)
  inference/                  unchanged (eval-time contiguous-KV engine; candidate to be replaced by PagedEngine later)
  envs/                       NEW  base.py harness.py math.py lean/{env.py,prompts.py} code.py
  train/                      engine: trainer.py train_module.py fsdp.py probes.py (+AdapterProbeView) sharded_ckpt.py ckpt_format.py
                              moe_module.py stage_gates.py stage_smoke.py   ([A2] the last two stay)
  algorithms/                 NEW  base.py registry.py
                              sft.py  dpo/  ppo/  grpo/  cispo.py  dapo.py  expert_iteration.py
  eval/                       unchanged + outcome_agreement.py (M7 only: see 4.5, changes eval_source_hash)
  recipes/                    NEW  sft.py dpo.py grpo.py ppo.py expert_iteration.py
  cli.py                      NEW  python -m align {sft,dpo,grpo,ppo,ei,eval}
```

Not moved on purpose: `config.py`, `exceptions.py`, `exit_codes.py`, `registry.py` (55 test import sites for `config.py` alone, all in
the freeze closure, and nothing is gained by a `core/` directory that the layer table does not already give); `model/`, `kernels/`,
`tokenizer/`, `rollout/`, `inference/`, `data/` names (97 test imports use `align.model` alone). Tests (121 files) stay flat in `tests/`; a
new `tests/contract/` holds the contract suites (2.3).

### 2.3 The algorithm contract

Today the pieces of an algorithm are spread over four places with no common name: a loss adapter `(model, batch) -> LossOutput`
(`train/train_module.py` `LossAdapter`), a stage config parsed by its own loader (`rl/stage_config.py:181`,
`ppo/stage_config.py:313`, `train/dpo_stage.py:285`), a launch-refusal function (`assert_a3_conformant` at
`rl/grpo_loss.py:337`, `assert_p4_cispo_conformant` at `rl/cispo.py:295`), and a `TrainModule` subclass. The contract names them.

```python
# align/protocols.py  (sketch)
@dataclass(frozen=True)
class Needs:
    reference: Literal["none", "adapters_off", "cached_logprobs", "model"]
    rollouts: bool
    critic: bool
class Algorithm(Protocol):
    name: ClassVar[str]
    Config: ClassVar[type]            # frozen dataclass, struct-mode loadable by align.config
    Batch: ClassVar[type]             # batch contract in align/batch.py, with .validate()
    needs: ClassVar[Needs]
    def __init__(self, cfg: "Config"): ...
    def validate(self, run: RunConfig, world_size: int) -> list[str]:
        """Launch refusals. Empty list = may launch. Replaces assert_a3_conformant et al."""
    def loss_adapter(self) -> LossAdapter:                      # (model, batch) -> LossOutput
    def train_module(self, model, optimizer, run: RunConfig) -> TrainModule: ...
    def make_source(self, data_cfg, *, env=None, engine=None, model=None, tokenizer=None): ...
    def extra_state(self) -> dict: ...                          # rides resume.pt["extra_state"] (exists)
    def load_extra_state(self, state: dict) -> None: ...
    def banner(self) -> dict:                                   # keys registered with align.registry
```

Rules the contract encodes, each already true somewhere in align and now stated once: `train_batch` is exactly
`(loss / grad_accum).backward()`; per-token tensors are label-aligned `[B, T-1]`; extra state loads with `weights_only=True`;
undefined metrics are NaN, never 0.0.

Mapping of the seven algorithms (current locations [read]):

| Algorithm | Loss today | Batch contract today | Config today | Extra state | Needs |
|---|---|---|---|---|---|
| `sft` | `make_sft_loss_adapter` `train/sft_loss.py:54` over `fused_linear_cross_entropy` (`kernels/fused_ce.py:259`) | dict from `data.loader.collate_sft` | `RunConfig` sections only | none | none |
| `dpo` | `make_dpo_loss_adapter` `train/dpo_loss.py:115` over `fused_dpo_loss` (`kernels/fused_dpo.py:545`) | `DpoBatch` `train/dpo_stage.py:818` | `DpoStageConfig` `:118` | tripwire response (`:680`) | `cached_logprobs` (`train/ref_cache.py`) |
| `ppo` | `ppo_loss` `ppo/loss.py:83`, `value_warmup_loss` `:163`, `PpoTrainModule` `ppo/trainer.py` | `PPORollout` `ppo/rollout.py:40` | `PPOStageConfig` `ppo/stage_config.py:110` | ratio tripwire (`ppo/health.py`) | `adapters_off` reference, critic, rollouts |
| `grpo` | `compute_grpo_loss` `rl/grpo_loss.py:179`, `GRPOTrainModule` `rl/grpo_module.py:53` | `RolloutBatch` `rl/types.py:21`, `TokenAdvRolloutBatch` `rl/grpo_loss.py:301` | `GrpoStageConfig` `rl/stage_config.py:92`, `GRPOLossConfig` `rl/grpo_loss.py:47` | `CollapseTripwire` state, `GrpoBatchSource` cursor (see "Still open" in DEPENDENCY_POLICY) | rollouts; reference only if `beta > 0` |
| `dapo` | `GRPOLossType.dapo` branch (`grpo_loss.py:194`) + dead-group filter (`rl/advantages.py:50`) + soft overlong (`rl/rewards.py:22`) | as `grpo` | as `grpo` | as `grpo` | as `grpo` |
| `cispo` | `cispo_policy_loss` `rl/cispo.py:110`, `fp32_head_forward` `:255` | as `grpo` | `GRPOLossConfig(loss_fn=cispo)` + `assert_p4_cispo_conformant` | as `grpo` | as `grpo`, `beta = 0` |
| `expert_iteration` | inner loss is `sft`; outer is an 11-stage driver `scripts/p2/expert_iteration.py` (850 lines) with chain-hash markers | stage markers (`RUN/round_NN/stages/<stage>.done.json`) | argparse in the script | round manifests | none (subprocess stages) |

**Two honest wrinkles in this table.**

- *"DAPO" in align is not DAPO's objective.* `GRPOLossType.dapo` is the clip-higher surrogate (DAPO Sec. 3.1, arXiv:2503.14476,
  `eps_high = 0.272` in `grpo_loss.py:54`), but the A3 gate refuses DAPO's own token-mean denominator
  (`grpo_loss.py:56` comment "A3 requires a float (Dr.GRPO)"; refusal text at `:356`, "not the DAPO token-mean 'token' default") in favour of a constant denominator
  and centred advantages without std division (Dr. GRPO, arXiv:2503.20783, Sec. 3.1 on the length and std normalisation biases).
  Dynamic sampling (DAPO Sec. 3.2) is the dead-group filter; soft overlong punishment (DAPO Sec. 3.4) is `rl/rewards.py:22`.
  The presets should say what they are: `dapo` (upstream-faithful, token mean, for comparison), `a3` (the pre-registered
  clip-higher + constant denominator), `cispo_p4` (CISPO, arXiv:2506.13585 Sec. 3.1 Eq. 3-5, no lower IS bound, `beta = 0`,
  fp32 head per Sec. 3.2). The A3 and P4 gates stay as `validate()` and keep refusing.
- *`expert_iteration` is an outer loop, not a per-step loss.* The contract has two kinds: a `LossAlgorithm` (everything above
  except this one) and an `OuterLoop` (a stage graph whose stages are subprocesses of other recipes). Do not force it into
  `loss_adapter`. It implements `validate`, `banner` and a `run(stage_graph)`; its stages call `recipes.sft`, `recipes.eval`.

**Contract test** (`tests/contract/test_algorithm_contract.py`, parametrised over the registry): (1) `Config` round-trips through
`dumps_toml` and `load_config` with struct mode; (2) `loss_adapter()` returns a finite `LossOutput` on a tiny fp64 model and
synthetic batch; (3) the gradient equals an eager fp64 reference to a stated tolerance (the "oracle tier" pattern already used
for model parity); (4) `extra_state()` survives `torch.save` / `torch.load(weights_only=True)`; (5) `Batch.validate()` rejects an
off-by-one label alignment; (6) `validate()` fires on each known-bad config in a table; (7) every `Config` field has a doc
comment (A6 above).

### 2.4 The environment contract

```python
# align/envs/base.py  (sketch)
GenerateFn = Callable[[list[GenRequest]], list[Generation]]    # the rollout-engine seam; exists in lean_multiturn.py:152,160
class Env(Protocol):
    name: ClassVar[str]
    def specs(self, split: str) -> Sequence[EpisodeSpec]: ...
    def run_group(self, specs: Sequence[EpisodeSpec], generate: GenerateFn) -> list[EpisodeRecord]:
        """Run all episodes of one prompt-group in lock step so the engine can batch across them."""
    def rows(self, ep: EpisodeRecord) -> list[TrainingRow]:
        """response_mask is 1 only on tokens the generator returned; turn_ids labels each token."""
    def reward_components(self) -> tuple[str, ...]: ...
```

- `MathEnv` (single turn): extract answer, `OutcomeVerifier`, `CompositeReward`; absorbs `rl/assembly.py:57` `assemble_rollout`'s
  scoring half.
- `LeanEnv(max_turns=K)`: `rl/envs/lean_multiturn.py` unchanged in behaviour; `max_turns=1` is the single-turn Lean prover.
  Layout `per_turn` or `concatenated` stays an env option; `PrefixBreak` stays.
- `CodeEnv`: new. Run candidate code against tests in a subprocess with the same constraints as `verify/lean/sandbox.py`
  (own session, `prlimit`, minimal environment, memory watchdog). Counts as a heavy job for the host preflight when pooled.

`GrpoBatchSource` becomes `RolloutSource(env, engine, algorithm_assembler)`; `MultiTurnBatchSource(GrpoBatchSource)` (a subclass
whose `bind` is "deliberately narrower than the parent's", `rl/multiturn_source.py:418`) disappears, because the parent no longer
assumes `problem/answer` rows.

### 2.5 Exact mapping, current to new

`F` = in the import closure of the running P2 scripts (4.1): do not edit or move until the chain exits. For an `F` file the new
path is created first as a facade that imports from the old path; the physical move and shim-direction flip happen in M7.
Files not marked `F` can move in M3 to M5 with an identity shim.

| Current | New | F |
|---|---|---|
| `align/__init__.py` | stays; new lazy names in M7 | F |
| `align/config.py`, `exceptions.py`, `exit_codes.py`, `registry.py` | stay | F |
| (new) | `align/protocols.py`, `align/batch.py` | |
| `utils/dist.py`, `utils/guard.py`, `utils/tracking.py` | `runtime/dist.py`, `runtime/guard.py`, `runtime/tracking.py` | F |
| `utils/host_state_gate.py`, `utils/supervisor.py` | `runtime/host_state_gate.py`, `runtime/supervisor.py` | |
| `utils/__init__.py` | stays (shim package) | F |
| `kernels/*` | stay; new `fused_policy_loss.py`, `paged_attention.py` | mixed |
| `tokenizer/*` | stay | F |
| `model/{attention,checkpoint,hf_config,mistral,moe,moe_ep,olmoe,olmoe_checkpoint,qwen3,qwen3_checkpoint}.py` | stay | F |
| `model/{gates,weight_average}.py` | stay | |
| `model/tp.py` **[A2]** | stay | F (imported by `tests/test_tp_bench.py`) |
| `ppo/lora.py` | `model/lora.py` | F |
| `ppo/value_head.py` | `model/heads.py` | F |
| `train/qwen3_stage.py` (`build_qwen3_stage_model`, `resolve_qwen3`, `validate_qwen3_cfg`, `qwen3_trainer_extras`, `qwen3_entry_banner`) | `model/build.py` (builders) and `recipes/sft.py` (banner, extras) | F |
| `train/qwen3_stage.py` `LoRASFTModule` | `algorithms/sft.py` (`LoRASFTModule`) | F |
| `train/qwen3_stage.py` `Qwen3Format` | `train/ckpt_format.py` (next to `DenseFormat`, `OlmoeFormat`) | F |
| `train/{adamw8bit,quant8,offload_adamw}.py` | `optim/{adamw8bit,quant8,offload_adamw}.py`; `make_optimizer` to `optim/factory.py` | F |
| `train/{trainer,train_module,fsdp,probes,sharded_ckpt,ckpt_format,moe_module}.py` | stay (the engine); `probes.py` gains `AdapterProbeView` from `ppo/trainer.py` | F |
| `train/stage_gates.py`, `train/stage_smoke.py` | **[A2] stay** (they import only `eval`, which is L2b; moving them to `recipes/` made `dpo_stage` import L5 and touched 12 frozen import sites for nothing). `recipes/` imports them | F |
| `train/__init__.py`, `eval/__init__.py`, `model/__init__.py`, `rollout/__init__.py`, `kernels/__init__.py`, `verify/__init__.py`, `tokenizer/__init__.py`, `data/__init__.py`, `inference/__init__.py` **[A2]** | stay | F |
| `ppo/__init__.py`, `rl/__init__.py`, `rl/envs/__init__.py` **[A2: had no destination]** | stay as legacy alias packages. `ppo/__init__` and `rl/__init__` import 9 and 10 submodules eagerly (names in `__all__`), so after M7 each re-export line goes through a shim: they are rewritten as lazy `__getattr__` re-exports to the new locations, in M7, with a test that `from align.rl import GRPOTrainModule` still works and that `import align.rl` imports no more than it does today | F (`rl/envs/__init__` free) |
| `train/sft_loss.py` | `algorithms/sft.py` | F |
| `train/dpo_loss.py` | `algorithms/dpo/loss.py` | |
| `train/dpo_stage.py` | `algorithms/dpo/{config,gates,source,adapter}.py`; `DpoBatch` to `batch.py` | |
| `train/ref_cache.py` | `algorithms/dpo/ref_cache.py` | |
| `rl/types.py`, `TokenAdvRolloutBatch` from `rl/grpo_loss.py` | `batch.py` | F |
| `rl/grpo_loss.py` | `algorithms/grpo/loss.py` | F |
| `rl/grpo_module.py` | `algorithms/grpo/module.py` | F |
| `rl/advantages.py`, `rl/rewards.py`, `rl/ops.py` | `functional/{advantages,rewards,ops}.py` **[A2: were `algorithms/common/` and `algorithms/grpo/`, which put an L4 module under the L2 `envs/lean`]** | F |
| `rl/tripwire.py`, `rl/canary.py` | `algorithms/grpo/{tripwire,canary}.py` | F |
| `rl/cispo.py` | `algorithms/cispo.py` (loss and conformance); `fp32_head_*` stay as the reference path for `kernels/fused_policy_loss.py` | F |
| `rl/stage_config.py` | `algorithms/grpo/config.py` | |
| `rl/batch_source.py`, `rl/assembly.py`, `rl/format_instrument.py` | `algorithms/grpo/{source,assembly,format_instrument}.py` | |
| `rl/multiturn_source.py` | `algorithms/grpo/multiturn.py` | |
| `rl/envs/lean_multiturn.py` | `envs/lean/env.py`; prompts vendored to `envs/lean/prompts.py` | |
| `ppo/gae.py` | `functional/gae.py` | F |
| `ppo/{loss,rewards,rollout,health}.py` | `algorithms/ppo/{loss,rewards,rollout,health}.py` | F |
| `ppo/trainer.py` | split: `algorithms/ppo/module.py` (`PpoTrainModule`), `algorithms/ppo/source.py` (`PPOBatchSource`), `AdapterProbeView` to `train/probes.py` | F |
| `ppo/stage_config.py` | `algorithms/ppo/config.py` | |
| `scripts/p2/expert_iteration.py` | library form `algorithms/expert_iteration.py`; script stays and later becomes a wrapper | frozen |
| `eval/verifier.py` | `verify/outcome.py` (old path becomes a shim). **[A2] changes `eval_source_hash`, see 4.5** | F |
| `eval/{golden,harness,metrics,proof_stats,roundtrip}.py` | stay | F |
| `data/outcome_agreement.py` | `eval/outcome_agreement.py` (**[A2]** adds a file to `align/eval/`, so it changes `eval_source_hash` too: M7 only) | F |
| `data/*` (rest) | stay | mixed |
| `verify/{equiv,expr,extract,normalize}.py`, `verify/lean/*` | stay | F (lean frozen by instruction) |
| `rollout/protocol.py` | stays | F |
| `rollout/{engine,paged,tp_engine}.py` **[A2]** | stay; `paged.py` gets `TritonPagedAttention` in M6 | F (imported by `tests/test_tp_bench.py`, which the chain launches) |
| `rollout/{client,native}.py` | stay; new `rollout/backends.py` | |
| `inference/{engine,bucketed}.py`, `inference/receipts.py` | stay | F / free |
| `scripts/{sft,dpo,grpo,ppo}.py` | stay; `recipes/*` hold the same flows; scripts become 30-line wrappers after M7 | F (sft) |

---

## 3. Minimal public API

Ten call sites. Names marked `exists` are in the repo today; `new` are the facade layer from M1/M2. The point is that
`scripts/sft.py` (261 lines), `scripts/grpo.py` (412) and `scripts/ppo.py` (482) are each one use of these ten.

```python
import align
# 1. config (exists: align.config.load_config). The loader splits [algo], [env], [rollout] before struct-mode parsing, so
#    config.py itself is not edited during the freeze.
run = align.load_run("configs/p4_cispo_lean.toml", overrides=["optimizer.lr=1e-6"])        # new, wraps load_config
# 2. host guards, before torch (exists: align.utils.guard.apply_resource_guards)
align.runtime.apply_guards(cpu_threads=8, heavy=True, label="cispo-lean")
# 3. model + tokenizer, family chosen by run.model.arch through a registry (new; today scripts/sft.py::build_model
#    and train/qwen3_stage.py::build_qwen3_stage_model)
model, tok = align.model.build(run)
# 4. algorithm by name or preset (new; the classes wrap today's GRPOTrainModule / DpoStage / PpoTrainModule)
algo = align.algorithms.create(run.algo)            # run.algo = {name="grpo", preset="cispo_p4", eps_high=...}
refusals = algo.validate(run, world_size=3)         # launch refusals as data; recipes raise if non-empty
# 5. environment (new; wraps rl/envs/lean_multiturn.py)
env = align.envs.make("lean/multiturn", run.env, verifier=align.verify.lean.StrictVerifier(...))   # StrictVerifier exists
# 6. rollout engine (exists: align.rollout.client.make_rollout_client; new name and registry)
engine = align.rollout.make_engine(run.rollout, model=model, tokenizer=tok)          # "native" | "paged" | "tp" | "spool"
# 7. batch source: algorithm-specific, built from env + engine (new; GrpoBatchSource is today's equivalent)
source = algo.make_source(run.data, env=env, engine=engine, model=model, tokenizer=tok)
# 8. trainer: the existing Trainer, built from the algorithm (new classmethod over Trainer.__init__)
trainer = align.train.Trainer.from_algorithm(run, model, algo, source, run_dir=run.paths.run_dir)
trainer.resume_from_latest()                        # exists
summary = trainer.train()                           # exists; honours HALT codes and coordinated stop
# 9. one-liner for the common case (new)
summary = align.fit(run)                            # steps 2-8 in order, with the stage gates from train/stage_gates.py
# 10. kernels as a plain function, for people who only want the kernel (new)
loss, metrics = align.kernels.fused_policy_loss(
    hidden, lm_head_weight, labels, advantages, old_logprobs,
    objective="cispo", eps_high=4.0, beta=0.0, temperature=1.0, chunk_tokens=2048,
)
```

Evaluation is the eleventh and only other entry: `align.eval.run(model, tasks=["minif2f"], n=32, verifier=...)`, which wraps
`eval/harness.py::run_eval`.

CLI: `python -m align grpo configs/p4_cispo_lean.toml --smoke` calls `recipes.grpo.run(...)`, so the existing torchrun command
lines in the docs keep their shape.

What is deliberately **not** public: `Trainer` internals, `TrainModule.optim_step`, shard layouts, `_`-prefixed helpers. The lazy
`align/__init__.py` pattern is kept (`import align` must not import torch; `tests/test_public_api.py` pins it).

---

## 4. Migration plan

### 4.1 The freeze closure

The task rule is: live GPU jobs run about two days; do not edit `scripts/p2/*`, `extras/vllm_rollout/*`, `align/verify/lean/*`,
`align/utils/guard.py`, `align/train/trainer.py`. The P2 chain also launches new Python processes every round, and each one
imports `align.*` fresh. A half-written or moved module breaks the next round even if the file itself is not on the list.

So I computed the static transitive import closure of the P2 entry points (`scripts/p2/*.py`, `extras/vllm_rollout/*.py`,
`scripts/sft.py`, `scripts/qwen3_export_check.py`, 20 files), following every `import` including lazy ones and every parent
package `__init__` [read, AST script run for this document]. Result for those entry points alone: 101 of 129 `align/` modules.
It is a conservative over-approximation (a lazy import on a path a P2 round never executes still counts), which is the safe direction.

**[A2] The first draft stopped there, but the chain scripts launch more than the P2 entry points.** `.tmp/p2/r1/chain2_gpu_validation.sh`
and `chain_after_baseline.sh` run `tests/test_tp_bench.py`, `tests/test_moe_ep_bench.py`, `scripts/olmoe_sft_smoke.py` and
`scripts/ckpt_convert.py` as `torchrun` steps. Adding those four entry points pulls in four more modules
(`model/tp.py`, `rollout/engine.py`, `rollout/paged.py`, `rollout/tp_engine.py`). Re-running the same AST closure gives
**105 of 129 reachable, 24 free**. This matters because M6 edits `rollout/paged.py` and M6/M7 both touch the TP engine. Rule: the
entry-point list of the closure is every `python`/`torchrun` target named in any `.tmp/p2/r1/chain*.sh` that is still pending, and
the closure test (`test_frozen_surface.py`) takes that list from a checked-in manifest, not from this document.

The 24 modules **not** in the closure, which may be edited or moved during the freeze:

```
align/data:      pretrain_corpus.py stageb_sources.py
align/inference: receipts.py
align/kernels:   _fa2_reference_upload.py fused_dpo.py ops.py
align/model:     gates.py weight_average.py
align/ppo:       stage_config.py
align/rl:        assembly.py batch_source.py format_instrument.py multiturn_source.py stage_config.py
align/rl/envs:   __init__.py lean_multiturn.py
align/rollout:   client.py native.py
align/train:     dpo_loss.py dpo_stage.py ref_cache.py
align/utils:     host_state_gate.py supervisor.py
align/verify/lean: sandbox.py
```

The names the frozen scripts and files import today (the contract the shims must keep) [read, AST]:

```
align.config: load_config                      align.registry: render_banner
align.data: hub; .ingest_parquet: iter_parquet_rows; .loader: collate_sft; .store: ShardDataset, ShardWriter
align.eval: proof_stats; .verifier: Outcome    align.exit_codes: PreflightRefused, run_with_exit_codes
align.model: qwen3, qwen3_checkpoint; .checkpoint: load_hf_safetensors
align.model.qwen3_checkpoint: align_to_hf, hf_to_align, read_qwen3_config, save_hf_qwen3
align.rollout.protocol: RolloutResult, read_request, spool_dirs, touch_heartbeat, write_result
align.train.qwen3_stage: build_qwen3_stage_model, qwen3_entry_banner, qwen3_trainer_extras
align.train.sft_loss: make_sft_loss_adapter    align.train.trainer: Trainer
align.train.stage_gates: GateLedger, assemble_sft_banner, assert_build_is_current, assert_distinct_devices,
    assert_probe_resolved, eval_source_hash, evaluate_and_gate, parse_args, record_golden_from_report,
    require_golden_receipt, run_entry_gates
align.train.stage_smoke: build_stage_model, build_training_sampler
align.utils: guard; .dist: barrier, init_distributed; .guard: apply_resource_guards, preflight_host_check
align.verify.lean: ReplPool, StrictConfig, StrictVerifier, find_toolchain, strict;
    .repl_pool: ReplError, ReplLimits, ReplProtocolError; .strict: AUDIT_PRELUDE, REF_NAME, lex, parse_statement
```

Names added by the four chain-launched entry points **[A2]**:

```
align.model: tp, moe, moe_ep, olmoe, olmoe_checkpoint; .qwen3: Qwen3Config
align.rollout.engine: PagedEngineConfig, SamplingParams     align.rollout.paged: kv_bytes_per_token, plan_capacity
align.rollout.tp_engine: TPPagedEngine
align.train: fsdp, sharded_ckpt; .ckpt_format: OlmoeFormat; .moe_module: build_moe_train_module; .offload_adamw: make_optimizer
align.train.trainer: MANDATORY_METRICS, Trainer, checkpoint_problem
align.utils.dist: local_rank, rank, seed_everything, world_size; .guard: assert_paths_off_root, host_state; .tracking: Tracker
align.data.store: ShardDataset      align.tokenizer: Tokenizer      align.model.olmoe_checkpoint: read_olmoe_config
```

`tests/test_frozen_surface.py` (M0) re-derives this list with the same AST scan on every run and asserts each name resolves, so
a refactor that drops one fails in CI rather than in the middle of round 3. **[A2]** The scan's entry-point list is a checked-in
manifest (`tests/frozen_entrypoints.txt`) that must contain every Python target named in a pending `.tmp/p2/r1/chain*.sh`.

### 4.2 Shim mechanics (verified)

For a file that moves, the old path becomes:

```python
# align/rl/grpo_loss.py (after the move)
import sys
from align.algorithms.grpo import loss as _m
sys.modules[__name__] = _m
```

Checked with a throw-away package outside the repo [cpu-probe, deleted afterwards]: after `import old.impl`, `old.impl is new.impl`
is `True`; `from old.impl import _private_name, SomeClass` returns the same objects; `from old import impl` returns the same
module; `unittest.mock.patch("old.impl.f", ...)` changes the behaviour seen through the new path. So `isinstance` checks, private
imports and monkeypatches in the distinct `align.*` import paths the tests use all keep working. (No test uses
`mock.patch("align...")`, `monkeypatch.setattr("align...")` or `import_module("align...")` on a moved path today; 0 hits. **[A2]**
re-checked by grep.) Pickle compatibility is not a concern: align checkpoints use safetensors and `weights_only=True`.

**[A2] What the identity shim does NOT preserve (found by a second probe, same method).**

1. **`python -m old.path` silently does nothing and exits 0.** The runner imports the shim as `__main__`; the shim replaces
   `sys.modules["__main__"]` with the real module and the real module's `if __name__ == "__main__":` block never runs (its
   `__name__` is the new dotted path). Probe: `python -m pk.old.gate` printed nothing and exited 0, `python -m pk.new.gate`
   printed `REAL MAIN RAN` and exited 3. `align/utils/host_state_gate.py:31` documents `python -m align.utils.host_state_gate` and
   that module is the D-state safety gate (it must return `EXIT_PREFLIGHT_REFUSED` when the host is wedged); a silent exit 0 would
   turn the gate into a pass. **Rule 3:** a moved module that has an `if __name__ == "__main__"` block gets a shim that forwards
   it: `if __name__ == "__main__": raise SystemExit(_m.main())`. The manifest records which moved modules have a `main`, and
   `test_legacy_shims.py` runs `python -m <old>` in a subprocess for each and compares the exit code with `python -m <new>`.
   Today only `host_state_gate.py` among the moved modules has such a block (`model/weight_average.py` has a documented CLI but
   does not move).
2. **Tests that read a module's source text now read the shim.** `tests/test_dpo_loss.py:190` greps `align/train/dpo_loss.py`
   for `from_pretrained`, `load_state_dict`, `AutoModel` and `NmoMistralForCausalLM`; `tests/test_h2_review_items.py:675-679`
   parses `align/utils/host_state_gate.py` for non-stdlib imports and reads `align/utils/guard.py`. After the shim they still
   pass, vacuously, because a 3-line shim contains none of those strings. See 4.6.
3. **Partial-initialisation cycles.** `ppo/__init__.py` and `rl/__init__.py` import their submodules eagerly. An old-path shim
   executes `from align.algorithms... import ...`, which runs `algorithms/__init__` and `algorithms/<x>/__init__`; if any of
   those import `align.rl` / `align.ppo`, the import order decides whether it works. The layering test cannot see this (shims are
   exempt). `test_legacy_shims.py` therefore imports every old and every new path in a **fresh interpreter, new path first and
   old path first**, and asserts both orders succeed.

Rules 1 and 2: shims warn only under `ALIGN_WARN_LEGACY_IMPORTS=1` (a default warning would add stderr noise to every P2 round);
and `tests/test_legacy_shims.py` asserts, from a manifest, that `import_module(old) is import_module(new)` for every moved path.

For an `F` file the direction is reversed first: the new module is a facade that imports from the old one. When the chain exits
(M7) the real code moves and the old path becomes the shim above. The shim stays for at least one release because
`scripts/p2/*` and `extras/` import the old paths and are not ours to edit.

### 4.5 [A2] `eval_source_hash`: moving anything in or out of `align/eval/` invalidates every gate ledger

`train/stage_gates.py::eval_source_hash` (`:94-113`) hashes the relative path and bytes of every `.py` under `align/eval/`. A
`GateLedger` records that hash; on load, if the hash differs, `_load` deletes every gate in `EVAL_DEPENDENT_GATES` (the golden
regression among them) and `require_golden_receipt` then refuses to train (`stage_gates.py:739-753`). Three planned moves change the
hash: `eval/verifier.py` becoming a shim, `data/outcome_agreement.py` arriving in `eval/`, and the M8 path-hygiene edit of
`eval/golden.py`. The first draft treated these as free once the chain had exited. They are not: every run directory that is
resumed after M7 needs the golden regression re-run (a 7B checkpoint job; `scripts/eval.py --mode=golden`).

Plan: (a) none of the three edits happens before the P2 chain exits (already true for M7, now stated for M8); (b) M7 and M8 land
as **one hash epoch**, i.e. adjacent commits with a single golden re-run after the second; (c) `tests/test_layering.py` and
`test_legacy_shims.py` print the new `eval_source_hash` so the epoch boundary is visible in CI; (d) the alternative that avoids
the epoch keeps `align/eval/` byte-identical: leave `Outcome`/`OutcomeVerifier` in `eval/verifier.py`; cut `verify -> eval`
(`verify/lean/reward.py:116`) by typing the result against a small Protocol in `verify/`; and put `outcome_agreement.py` in
`recipes/` (L5, which may import `eval`) instead of `eval/`, with `align.data.__getattr__` serving the old name lazily. It costs a
slightly less tidy layout and saves a 7B golden re-run, so it is the cheaper choice unless the owner wants the tidier layout. The
owner decides (open question 11).

### 4.6 [A2] Tests that pin source text at old paths

Reading file text of a library module or of `scripts/*.py` is an established pattern here (the CPU suite cannot execute the GPU
scripts, so it pins them by structure). Found by grep: `test_dpo_resume_wiring.py` (string-matches `scripts/dpo.py`:
`'"--resume"'`, `os.environ.get("ALIGN_RESUME") == "1"`, the rotation guard), `test_grpo_resume.py:293` (AST of `scripts/grpo.py`:
`resume_from_latest` call ordering), `test_dpo_stage.py:1566` (`assert_launch_topology(cfg)` before `run_entry_gates(`),
`test_ref_cache.py:179,300` (`scripts/dpo.py`, `scripts/precompute_ref_logprobs.py`), `test_h2_review_items.py:374,383`
(`train()` in `scripts/sft.py`, `pretrain.py`), `test_qwen3_stage.py:301`, `test_g0_integration.py:38`, `test_olmoe_scripts.py`
(several), `test_verifier.py:371` (`scripts/eval.py`), plus the library-text pins in 4.2 item 2.

Consequences for the plan: (1) **M7 "make `scripts/*` 30-line wrappers" cannot keep the baseline green** unless each pin is
retargeted to the `recipes/` file that now holds the flow, with the same assertion, in the same commit. M2's recipes are copies,
so the pins should be written against both files from M2 on (a parametrised `[script, recipe]` pin), and the script pin is
dropped only at M7. (2) M3's move of `train/dpo_loss.py` (and `utils/host_state_gate.py`) must retarget the two library-text
pins to the new path in the same commit, and the shim manifest records the pin so `test_legacy_shims.py` fails if a pinned
path is a shim (a vacuous pin is a lost guard, not a pass). (3) The gate "same passed counts" cannot hold literally once
contract tests are added; it is restated in 4.3 as "every baseline test id still passes".
### 4.3 Steps
Every step has the same gate: `pytest -m "not gpu and not oracle and not slow"` where **every test id in the M0 junit baseline
still passes or is skipped for the same reason** (**[A2]** not "equal counts": M1 onward adds contract tests, so the totals must
grow, while a baseline id going missing or silently skipping must fail), `tests/test_dependency_firewall.py` green,
`tests/test_public_api.py` green (import stays cheap), `tests/test_frozen_surface.py` green, and, from M3, `tests/test_legacy_shims.py`
(identity, `-m` forwarding, both import orders, no vacuous source pin: 4.2 and 4.6). All GPU-free. All run with `CUDA_VISIBLE_DEVICES=`.
**[A2]** The suite shares a 48-core box with the live chain: run it under `nice -n 10`, `OMP_NUM_THREADS=4`, `-p no:cacheprovider`,
and only after the host preflight in `/apps/CLAUDE.md` (`free -g; uptime`) passes; `conftest.py` already keeps `tmp_path` off `/tmp`.
| Step | What | Effort | Touches `F` files? | Extra gate |
|---|---|---|---|---|
| **M0** | Safety net, no code moves: record the real CPU baseline (junit file); add `test_frozen_surface.py`, `test_legacy_shims.py` (empty manifest), `test_layering.py` (report-only, lists today's violations from 2.1) | 0.5 d | no | baseline numbers committed in the doc |
| **M1** | Facade layer, additive: `protocols.py`, `batch.py` (re-exporting the existing batch classes), `algorithms/{base,registry}.py` and `algorithms/{sft,dpo,ppo,grpo,cispo,dapo}.py` as classes that wrap the existing loss adapters and `TrainModule` subclasses; contract test suite 2.3 | 1.5 d | no (imports only) | contract tests pass for all six |
| **M2** | `recipes/` and `cli.py`: the `main()` flows of `scripts/{sft,dpo,grpo,ppo}.py` as functions, written as new files (the scripts are not edited). `align.load_run` splits `[algo]/[env]/[rollout]` before `load_config`. | 3 d | no | tiny-model parity: `recipes.sft` and `scripts/sft.py --smoke` produce identical `history.jsonl` losses on CPU/gloo (`train/stage_smoke.py` already builds the model); same for grpo and dpo smoke |
| **M3** | Move the non-`F` modules with identity shims: `dpo_*`, `ref_cache`, `rl/{stage_config,batch_source,assembly,format_instrument,multiturn_source}`, `ppo/stage_config`, `utils/{host_state_gate,supervisor}`. Fix cycles that need only free files. **[A2]** Retarget the two library-text pins (`test_dpo_loss.py:190`, `test_h2_review_items.py:675`) in the same commit; `host_state_gate` gets the `__main__`-forwarding shim | 2 d | no | `test_legacy_shims.py` manifest grows by exactly the moved paths; `python -m align.utils.host_state_gate` exit code equals the pre-move exit code |
| **M4** | `envs/`: `Env` protocol, `MathEnv`, `LeanEnv` from `rl/envs/lean_multiturn.py`; vendor `lean_prompts.py` as `envs/lean/prompts.py` with a sha256 identity test against `extras/vllm_rollout/lean_prompts.py`; extract `displayed_statement_matches`; collapse `MultiTurnBatchSource` into `RolloutSource(env=...)` | 3 d | no | `tests/test_lean_multiturn_env.py`, `test_multiturn_source.py`, `test_lean_gen_contract.py` unchanged and green; GRPO smoke bit-equal before and after |
| **M5** | Kernel A/B: `kernels/fused_policy_loss.py` as a chunked `autograd.Function` (5.4-A) with an eager fp64 oracle test; `mm(out_dtype=float32)` probe on GPU after the freeze | 3 d CPU + 0.5 d GPU | no | fp64 gradient oracle (max abs error recorded); GPU parity vs `fp32_head_logprobs` |
| **M6** | Kernel C: Triton paged-attention decode (5.4-C) behind `TritonPagedAttention`, `fast` and `invariant` modes | 8 d (CPU interpreter tests 3 d; GPU validation and tuning 5 d) | **[A2] yes: `rollout/paged.py` is in the closure of `tests/test_tp_bench.py`; starts only after the chain exits** | bit-invariance test: row output independent of batch composition; `tests/test_paged_engine*.py` green; decode benchmark vs `TorchPagedAttention` |
| **M7** | After the P2 chain has exited (gate **[A2]**: `pgrep -f '\.tmp/p2/r1/chain'` and `pgrep -f 'scripts/p2/'` both empty, not only `chain_after_baseline`: `chain4.sh` and `chain5.sh` are live today; `nvidia-smi` idle, as in `docs/P3_ENGINE.md` GPU plan): move the `F` files, flip each facade to a shim, rewrite the three legacy `__init__`s as lazy re-exports, add the new lazy names to `align/__init__.py`, retarget the script-text pins (4.6) and make `scripts/*` 30-line wrappers, sweep path mentions in `docs/` (README, `ARCHITECTURE.md`, `DEPENDENCY_POLICY.md` and 6 more name `align/rl/`, `align/ppo/`, `align/utils/`, `align/train/...`) | **3 d** (was 1.5: about 55 `F` modules, ten-odd pins, three `__init__`s, docs) | yes | full suite and frozen-surface test after **each** moved file, not at the end; then one golden re-run (4.5) |
| **M8** | Open-source hygiene: remove the hard-coded paths (2.1), `Nmo*` names (aliases kept; 38 files mention `nmo`/`NMO`), accept both `ALIGN_*` and `NMO_*` env vars, add a LICENSE (none exists in the repo root), relax `torch==2.11.0` / `triton==3.6.0` / `requires-python <3.13` pins to tested ranges, CI lane with ruff + firewall + layering, README rewritten around the ten call sites | 3 d | yes | clean checkout test: `uv venv && uv pip install -e .[dev]` and the CPU suite pass with no `/apps` path present |

Total about 27.5 working days (about 5.5 weeks for one engineer; **[A2]** M7 re-estimated from 1.5 to 3 d). M1 to M5 are all runnable
during the freeze because they only add files or move non-`F` ones. M6 (it edits the frozen `rollout/paged.py`) and the GPU halves
of M5 wait for free GPUs.

### 4.4 Risks

| Risk | Mitigation |
|---|---|
| A new file in an existing package changes import order or a `__init__` side effect | Only new subpackages are added during the freeze; no existing `__init__.py` is edited until M7 |
| The facade and the old class drift while both exist | Facades hold no logic; contract tests call both paths and compare |
| `sys.modules` shim hides a circular import | `test_layering.py` is strict before M7 flips anything |
| M2 recipes diverge from scripts that P2 still runs | Parity test on tiny model; the scripts are the source of truth until M7; the source pins of 4.6 are parametrised over `[script, recipe]` from M2 |
| Test counts in docs are stale, so "stay green" has no baseline | M0 records it |
| ~~Moving `train/stage_gates.py`~~ **[A2]** | No longer moved (2.1). `eval/` changes still invalidate the ledger, see 4.5 |
| **[A2]** `python -m <old.path>` through a `sys.modules` shim silently does nothing and exits 0 | `host_state_gate.py` is a safety gate with a documented `python -m align.utils.host_state_gate`. Shim rule 3 in 4.2 plus a subprocess test |
| **[A2]** Source-pinning tests read old file paths and script text (4.6) | Retarget each pin in the same commit as the move, and fail the gate if a pin goes vacuous |
| **[A2]** Eval source hash epoch (4.5) | Plan one golden re-run in M7 |

---

## 5. Where custom Triton kernels pay off

### 5.1 Method and constants

For each candidate the question is: **is the op bound by bytes or by FLOPs, and how much of the iteration is it?** A kernel can
only recover the gap between what the code does today and the roofline. Constants:

| Constant | Value | Source |
|---|---|---|
| A10G HBM bandwidth | 600 GB/s | **[assumed]** datasheet, same value used in `docs/P3_ENGINE.md` |
| bf16 tensor-core peak / best measured | about 70 / 68 TFLOP/s (one `8192x4096 @ 4096x4096`); 44 to 53 TFLOP/s in-layer | [read] `docs/MEASUREMENTS.md:102-105,254` |
| fp32 CUDA-core peak, sustained | 31.2 peak, 25 sustained (sensitivity 20 to 31) | **[assumed]** datasheet; not measured here. TF32 is disabled in `rl/cispo.py::_strict_fp32_matmul` |
| Qwen3-8B (Goedel-Prover-V2-8B) | hidden 4096, 36 layers, 32 q heads / 8 kv heads, head_dim 128, vocab 151,936, intermediate 12,288, untied head | [read] model `config.json` |
| KV bytes per token | 2 x 36 x 8 x 128 x 2 B = 147,456 B (144 KiB) | [calc]; matches `docs/P3_ENGINE.md` (3.60 GiB / 26,192 tokens) |
| Weights bf16 | 15.26 GiB = 16.39 GB; 27.3 ms to read once at 600 GB/s | [calc]; 15.26 GiB is `docs/P3_ENGINE.md` |
| Eager launch cost | 6 us per op | **[assumed]**, typical 5 to 10 us |

### 5.2 Where the time goes in one RL iteration (why decode dominates)

Illustration with stated knobs, not a measurement. 16 prompts x 8 samples, mean completion 4,096 tokens (cap 16,384): 524k
generated tokens per iteration.

- Generation at vLLM-class throughput (about 876 tok/s end to end, the vLLM figure quoted in `docs/P3_ENGINE.md`): about 600 s. At the
  pure-torch engine's modelled tp=4 rate (218 to 386 tok/s, rows 3 to 5 of the table in 5.4-C; **[A2]** the first draft's low end,
  100 tok/s, came from a tp=1 batch of 16 that does not fit in memory, see 5.4-C): 1,360 to 2,400 s.
- Training on the same tokens, LoRA with activation checkpointing: body forward `2 x 6.95e9 x tokens` (the 36 blocks hold 6.95 B
  parameters [calc]); with recompute and a dX-only backward, about 3x forward = 41.7 GFLOP/token; 524k tokens = 21.8 PFLOP;
  on 3 train GPUs at 45 TFLOP/s with no comm: about 160 s. FSDP over host shared memory is comm-bound on this box
  (`docs/MEASUREMENTS.md` section (h), 92% of a micro-step), so 160 s is a floor.

Reading: with the native engine, generation is about 8x to 15x the training body compute (the 160 s above excludes the head,
which adds about 35 s on 3 GPUs: 32 micro-batches x 3.3 s / 3, see 5.4-A); even against vLLM it is the larger term. The
decode path is where a kernel changes the iteration time; the head is a second-order term. This ordering drives 5.5.

### 5.3 What exists

- Triton: `kernels/flash_attention.py` (FA2 forward and backward, adapted from the user's `Flash-Attention-101/triton/flash_attention.py`,
  kept as `_fa2_reference_upload.py`), `kernels/ops.py` (RMSNorm/RoPE/SwiGLU, pedagogy, off by default). The FA2 backend is not
  selectable at HEAD: no gate artifact at the gate path (REFERENCE_GAP_ANALYSIS 1.1 item 9); SDPA is the working backend.
- **Not Triton, despite the name:** `kernels/fused_ce.py` and `kernels/fused_dpo.py` are chunked `torch.autograd.Function`s. They
  implement the Liger FusedLinearCrossEntropy design (arXiv:2410.10989, "FusedLinearCrossEntropy" paragraph, Eq. 15-17: chunk the
  head, write `dh = W dlogits` and `dW += h dlogits^T` during forward) in torch matmuls.
- `rollout/paged.py::TritonPagedAttention` is an interface with `raise NotImplementedError`. Its docstring already states the
  kernel contract, including bit-invariance ("fixed KV-split count per context length, no atomics").

### 5.4 Candidates

#### A. Fused head + policy-gradient token loss (with logprob gather)

**Today** (`rl/cispo.py::fp32_head_logprobs`, `:202`): per chunk of 2,048 tokens, `h.float() @ W32.T` in strict fp32, `logsumexp`,
`gather`; `torch.utils.checkpoint` per chunk, so backward recomputes the GEMM; a separate no-grad pass for entropy.

**GEMM count** at `2*H*V = 1.245 GFLOP/token` (20.4 TFLOP per 16,384-token micro-batch [calc]), frozen head (LoRA):

| Path | GEMMs | Arithmetic | Time per 16k tokens |
|---|---|---|---|
| Today: fwd + checkpoint recompute + `dh` + entropy pass | 4, strict fp32 | 4 x 20.4 / 25 | **3.3 s** (2.6 to 4.1 for 31 to 20 TFLOP/s) |
| Fused: fwd with `dlogits` formed in-chunk, then `dh = dlogits @ W`; entropy from the same sweep | 2, bf16 operands, fp32 output | 2 x 20.4 / 50 | **0.82 s** (0.68 to 0.91 for 60 to 45) |

Ratio about 4x [calc, with the assumed fp32 rate]. If the head is trainable (full FT), add one `dW` GEMM to each side
(5 vs 3): the ratio is the same.

**Why the in-chunk form works for this loss.** For CISPO the per-token gradient **of the objective J** with respect to the logit vector is
`c_t (softmax(z) - onehot(y))` with `c_t = -A_t * sg(clip(r_t, ., 1+eps_high))` (CISPO Eq. 4-5, arXiv:2506.13585 Sec. 3.1;
the weight is a stop-gradient constant). **[A2]** The sign is that of the ascent direction: the loss `-J` that the optimizer
minimises has `c_t = +A_t * sg(clip(r_t, ., 1+eps_high))`. Checked in fp64 autograd on a 11-way softmax: both signs reproduce
the autograd gradient to 7e-18, and the two differ only by the overall sign. The kernel must use the loss sign. `r_t` needs the new logprob, which the same chunk has just computed, so `c_t` is known
before `dlogits` is formed. The KL term adds a closed-form function of the same logprob. This is exactly the structure of
`fused_ce._FusedLinearCrossEntropy` (gradient written in forward, backward is a scalar rescale), extended with the token
coefficient `c_t`. For the clipped-surrogate objectives (DAPO/A3) `c_t = -A_t r_t 1[not clipped]`, also known in-chunk.

**The precision insight (the real reason the fp32 head matters).** MiniMax (arXiv:2506.13585 Sec. 3.2) attribute the
train/inference probability mismatch to "high-magnitude activations in the LM head" and fix it with an FP32 head (correlation
about 0.987 to 0.997, their Fig. 3). In align the head weights and the final hidden state are bf16. A bf16 product is exact in
fp32 (8+8 mantissa bits < 24), and both a bf16 tensor-core GEMM and an fp32 GEMM accumulate in fp32. So an "fp32 head" built
from bf16 operands differs from a bf16 GEMM **only in whether the output logits are rounded to bf16**. At `|z| ~ 16..32` the bf16
spacing is 0.125, so a logit is off by up to 0.06 and so is a logprob.

[cpu-probe] N=512, H=1024, V=16,384, random bf16 `h` and `W` with logits up to `|z| = 50.6`; fp64 reference of the same bf16
operands. fp32-head max abs logprob error **7.1e-6**; bf16-output matmul max abs logprob error **0.145** (mean 0.049). That is
the whole effect, reproduced on CPU, with operands identical in both arms.

Consequence: the head does not need fp32 *compute*; it needs fp32 *output*. `torch.mm(a_bf16, b_bf16, out_dtype=torch.float32)`
exists in torch 2.11 (the op `aten::mm.dtype` is registered; on CPU it reports "could not run on the CPU backend", so it is
CUDA-only and **unverified on a GPU here**). The current `fused_ce.py:` line `logits = torch.matmul(h_chunk, weight.t()).to(cdtype)`
rounds to bf16 first and upcasts after, so the SFT loss carries this rounding noise too (harmless for SFT gradients, which are
unbiased in aggregate; it matters for ratios). The same trick makes the P4 fp32 head about 3x to 4x cheaper than strict-fp32
and also saves the fp32 weight copy (`weight.float()` is 2.49 GB, against 1.25 GB bf16 [calc]) on a 23 GB card.

Caveat to test on GPU: tensor-core accumulation is not IEEE fp32 (reduced alignment and truncation inside the MMA). The size of
that error against bf16 output rounding must be measured with the probe above run on CUDA.

**[A2] The backward GEMM is not "free precision" either.** The second fused GEMM, `dh = dlogits @ W`, only runs at the bf16
tensor-core rate if `dlogits` is cast to bf16 first (a strict-fp32 `dlogits @ W` is exactly the cost being removed). The forward
logprob keeps the 7e-6 accuracy above, but `dh` does not. [cpu-probe] N=512, H=1024, V=16,384, random bf16 `h`, `W`, token
coefficients `c_t ~ N(0,1)`: relative L2 error of `dh` against fp64 is **3e-7** with an fp32 `dlogits` and **1.8e-3** with a
bf16-rounded `dlogits` (max abs 1.8e-3 against `max|dh|` 0.68). That is ordinary bf16 backward noise, but it means the fused
path is not bit-comparable with the strict-fp32 path, the M5 oracle test needs a tolerance of about 5e-3 relative on `dh` (and
1e-4 on logprobs), and "fp32 head" in the P4 pre-registration should be read as "fp32 logits", not "fp32 gradient". The cast adds one
read and one write of the tile (about 1.9 GB per 2,048-row chunk, 25 ms per 16k tokens), already inside the 10% to 16% range below.

**Share of a training step.** Body about 3 x 228 TFLOP at 45 TFLOP/s is about 15 s per 16,384-token micro-batch [calc]; head 3.3 s
today versus 0.82 s fused: a saving of 2.5 s of about 18.3 s, **about 13% of the training half of an RL iteration**, and about
**[A2]** 1% to 3% of the iteration end to end (32 micro-batches x 2.5 s / 3 GPUs = 27 s against an iteration of about 1,550 to 2,600 s
with the native engine, about 800 s against vLLM; 5.2). It is worth doing because it is cheap and also removes 1.2 GB of resident memory, not because
it moves the iteration.

**Is Triton needed?** No. Elementwise traffic per 2,048-row chunk: fp32 logits are 1.245 GB; today about 8 passes (GEMM write,
`logsumexp` read, recompute write and read, softmax write and read, `dlogits` read) is 10 GB, or 16.6 ms at 600 GB/s,
133 ms per 16k tokens [calc]. Against today's 3.3 s that is 4%. **[A2] The denominator that matters for the Triton decision is
the fused path's, not today's.** After A the head costs 0.82 s, and its elementwise tail is about 7 passes (logits write,
`logsumexp` read, softmax read, `dlogits` write, bf16 cast read and write, `dh` GEMM read, about 7.5 GB per chunk, 100 ms per 16k
tokens): **10% to 16% of the fused head**, 0.5% to 0.7% of the 18.3 s training step. A 3x faster fused epilogue (the Liger
figure below) would recover about two thirds of that, at most 0.07 s per micro-batch. The conclusion (do not write Triton for it)
stands, but at the stated 4% the reasoning understated the residual by 3x. Fusing it (Liger reports a 3x faster *cross-entropy kernel* in isolation,
arXiv:2410.10989 Fig. 2a and 3a, at vocab 163,840, benchmarking the kernel alone with the GEMM excluded) would recover at most about 0.1 s per micro-batch. AgileRL's
own note that its Liger GRPO kernel gives no speedup (1.2) is the same fact. Ship it as a chunked `autograd.Function`
(`kernels/fused_policy_loss.py`), tested against `fp32_head_logprobs` and an fp64 oracle.

#### B. fp32-head `logsumexp` as a standalone kernel

Does not pay. After A it is an epilogue over a `[chunk, V]` fp32 tile, about 10% to 16% of the fused head's time (4% of today's), and a Triton GEMM on sm_86
(99 KB shared memory, `docs/MEASUREMENTS.md:124`) is unlikely to beat cuBLAS on these shapes. The same document reached the same
conclusion for the grouped-GEMM kernel (P1-12 fails its own entry condition). The only reason to write it later is to avoid
materialising the `[chunk, V]` tile at all (memory, not time), and chunk size already bounds that: at 512 rows the tile is 311 MB.
#### C. Paged-attention decode (the one that pays)
**Today** (`rollout/paged.py::TorchPagedAttention`, `'segment'` mode, read above): for every decode segment, every layer:
`cache.gather` copies K and V into a contiguous strip (read L + write L), a second copy into a zero-filled padded buffer
(zero fill `L_max`, copy read L + write L), then one SDPA over `L_max` with a boolean mask. Per sequence per layer per tensor
(K or V) the traffic is about `4L + 2 L_max`, against `L` for a fused kernel that reads blocks once through the block table.
With `rho = L_max / L_mean`, the multiplier is `4 + 2 rho`: 8x at `rho = 2`. The length distribution of a prover RL batch is
heavy-tailed (most stop in a few thousand tokens, a few reach the 16k to 40k cap, `rollout/paged.py` docstring), so `rho >= 2` is
the normal case, and `docs/P3_ENGINE.md` already flags the padded transient (1.17 GiB per layer call at 64 sequences, 4.7 GiB if SDPA
falls back to the math backend with GQA expansion).

On top of bytes there is launch overhead: about 12 eager ops per segment per layer is `B x 36 x 12` launches per step
(27,648 at B = 64), about 166 ms at 6 us [assumed]; the GPU is idle while Python issues them.

Model (all [calc] from 5.1; `tp` as in `docs/P3_ENGINE.md`; comm term is that doc's assumed 20 to 26 ms for B = 64 at tp = 4 and
about 10 ms for B = 16; fused kernel at 85% of HBM peak [assumed]; step = max(GPU time, CPU launch time) for today):

| Setting | KV tokens (capacity) | Weights ms | KV ideal ms | Today step ms (tok/s) | Fused step ms (tok/s) | Speedup |
|---|---|---|---|---|---|---|
| tp=1, B=6, ctx 4,096, rho=2 | 24,576 (26,192) | 27.3 | 6.0 | 76 (79) | 34 (174) | 2.2x |
| tp=4, B=16, ctx 4,096, rho=2 | 65,536 (438,048) | 6.8 | 4.0 | 49 (326) | 22 (743) | 2.3x |
| tp=4, B=64, ctx 4,096, rho=2 | 262,144 (438,048) | 6.8 | 16.1 | 166 (386) | 52 (1,240) | 3.2x |
| tp=4, B=48, ctx 8,192, rho=2 | 393,216 (438,048) | 6.8 | 24.2 | 220 (218) | 55 (870) | 4.0x |

**[A2] Rows removed from the first draft because they do not fit in the card.** The first draft listed tp=1 B=16 at ctx 4,096 and
8,192 (3.4x and 4.4x) and tp=4 B=64 at ctx 8,192 (4.2x). The KV pool the native engine can allocate is in `docs/P3_ENGINE.md`
("Native engine, bf16 KV" table): 26,192 tokens at tp=1 (weights 15.26 GiB leave 3.60 GiB), 163,472 at tp=2, 438,048 at tp=4 (util
0.90). tp=1 B=16 at ctx 4,096 needs 65,536 tokens (2.5x over) and at 8,192 needs 131,072 (5x over); tp=4 B=64 at 8,192 needs
524,288 (1.2x over). Those speedups were computed correctly for the stated inputs but describe a batch the engine cannot hold, and
the 3.4x to 4.4x figures were the top of the quoted range. Replacement rows use the largest batch that fits (B=6 at tp=1, B=48 at
tp=4 with ctx 8,192; 20 ms comm for B=48). Row arithmetic, checked: KV ideal = tokens x 147,456 B / tp / 600 GB/s; today = weights
+ (4+2rho) x KV ideal + comm (launch-bound at B=64: 27,648 x 6 us = 166 ms); fused = weights + KV ideal / 0.85 + comm.

Expected: **2x to 4x on decode steps at the largest batch that fits (6 to 64) and rho about 2**; larger with a larger `rho` or longer context; at batch 1
the step is weight-bound and the kernel changes little (it still removes the launch overhead, 2.6 ms vs 27 ms). Amdahl: this
speeds the generation half, which 5.2 shows is the larger term. The bar in `docs/P3_ENGINE.md` (native at least 0.5x vLLM
throughput) is the acceptance test, and the stated fallback there ("the Triton kernel is the next step, not TP") is consistent
with this analysis. The vLLM kernels in the PagedAttention paper (arXiv:2309.06180 Sec. 5.1, "Kernel-level Optimization") are three fusions:
fused reshape-and-block-write ("to minimize kernel launch overheads"), fused block read and attention (adapted from the FasterTransformer
attention kernel, reading the KV cache through the block table), and fused block copy.

**Design constraint from RL, not performance:** the per-row key reduction must not depend on batch composition (the
`invariant` mode of `TorchPagedAttention`, `rollout/paged.py`). Use a split-KV (flash-decoding style) kernel with a **fixed split
size per context length** (e.g. 512 keys per split, so the split count is `ceil(L / 512)`, a function of `L` only), a second
reduce kernel in a fixed order, no atomics. That keeps `rho` (the train/infer logprob ratio) unaffected by who shares the batch.
It is the user's FA2 forward loop (`Flash-Attention-101/triton/flash_attention.py`, online softmax per KV block) with `Lq = 1`,
the KV loop split across programs, and the block table as an extra indirection; the numbered-note house style applies
(one note on split-KV online softmax merging, one on block-table indexing, each with a notebook that checks the kernel against
the naive version).

**Validation plan (needs free GPUs):** (1) bitwise: same sequence, run alone and inside batches of 2, 16, 64 with different
neighbours: identical logits. (2) numeric: against `TorchPagedAttention(mode="row")` at bf16, max abs logit error stated.
(3) throughput: `tests/test_paged_engine_bench.py --gpu` with the same prompt ids as the vLLM manifest. (4) peak memory with
`--reserve-gib` reduced, since the padded transient goes away.

#### D. Candidates that do not pay

| Candidate | Estimate | Verdict |
|---|---|---|
| Per-token CISPO/GRPO loss arithmetic alone | tensors of `B x T` fp32: 64 KB at 16k tokens; about 10 elementwise ops is below 1 MB of traffic, about 50 us of launch time against a multi-second step | Do not write a kernel |
| `adamw8bit` Triton step (open item in DEPENDENCY_POLICY) | ideal about 18 GB of traffic for 1.8 B local elements [calc, about 10 B/element] = 30 ms; runs once per optimizer step, and a step is 8 to 85 micro-steps of seconds each | Under 1% of a step. No |
| Fused LoRA add (`W x + B A x`) | one saved read+write of each `[T, out]` bf16 activation per projection: 2.55 GB per layer at 16k tokens = 4.3 ms, 153 ms per forward over 36 layers [calc] | About 1% of a 15 s body pass; no |
| FA2 Triton as default | Gate exists and is not enabled (no artifact at the gate path); attention is about 4% of layer FLOPs at seq 1k (`kernels/fused_ce.py` docstring), more at 16k | Finish the equivalence gate first; a measurement task, not a kernel task |
| RMSNorm/RoPE/SwiGLU Triton (`kernels/ops.py`) | kept flag-gated and off; `torch.compile` fuses them | Keep as pedagogy |

### 5.5 Ranking

1. **C, paged-attention decode.** About 2x to 4x on decode (3x central), on the largest term. 8 days including GPU tuning. Needs free GPUs.
2. **A, fused head + policy loss as an `autograd.Function` with `out_dtype=float32`.** About 4x on the head, about 13% of the
   training half, 1.2 GB less resident memory, and it fixes the SFT output-rounding as a side effect. 3 days plus 0.5 day GPU.
   No Triton.
3. Everything else in 5.4-D: no.
---
## 6. Open questions for the owner
1. **Name.** `align` is a common PyPI name. An open-source release needs a distinct distribution name; the import name can stay.
   Renaming `Nmo*` (38 files) is mechanical; do you want the old names kept as aliases for one release?
2. **Licence.** The repo has no LICENSE. AgileRL is Apache-2.0, Liger-Kernel is BSD-2-Clause, open-instruct (vendored GRPO loss,
   `rl/grpo_loss.py`) is Apache-2.0: Apache-2.0 for align is compatible. Confirm.
   **[A2]** The vendored open-instruct code (`rl/__init__.py` says "vendored verbatim") also needs its upstream copyright
   notices kept in the files and a `NOTICE` entry before release; M8 does not list that, add it.
3. **Is `expert_iteration` an algorithm or a recipe for you?** The proposal makes it an `OuterLoop` with its own small interface
   (2.3). If you would rather keep it a script, drop it from the registry and keep only `recipes.sft` and `recipes.eval`.
4. **A10G constants.** Section 5 rests on 600 GB/s and a 25 TFLOP/s sustained fp32 rate, both assumed, plus 6 us eager launch
   cost. Three 2-minute microbenchmarks (a bandwidth copy, an fp32 and a bf16 GEMM at the head shape, an empty-op loop) on a free
   GPU would turn every [assumed] in 5.1 into a measurement. They are GPU jobs and must take the heavy-job lock.
5. **`mm(out_dtype=float32)` on CUDA.** Not runnable here. If it is absent or slow for the head shape the fallback is a Triton
   GEMM with fp32 epilogue, which moves A from 3 days to about 6 and changes the 5.5 ranking only slightly.
6. **Tensor-core accumulation error versus bf16 output rounding** for the head (5.4-A caveat): needs the CPU probe repeated on a GPU.
7. **Test layout.** AgileRL mirrors the package in `tests/`. This proposal keeps the 121 flat files and adds `tests/contract/`.
   Mirroring would change documented `pytest tests/test_x.py` commands in several docs; say if you want it.
8. **Which `[algo]` name for GSPO-style trajectory ratios.** Not implemented in align. Add it in M1 or leave it out until a run needs it?
9. **`inference/engine.py` (contiguous KV) versus `PagedEngine`.** After M6 the paged engine can serve eval too, and
   `inference/` could be deleted (962 + 202 + 220 lines). Not proposed now because eval receipts and golden regression depend on its exact numerics.
10. **Test baseline.** README says 1608 passed, DEPENDENCY_POLICY 1089, REFERENCE_GAP_ANALYSIS 1107. M0 resolves it; until then no
    "stays green" claim here has a number behind it.
11. **[A2] Eval hash epoch.** Moving `Outcome` into `verify/` and `outcome_agreement.py` into `eval/`, and the M8 edit of
    `eval/golden.py`, change `eval_source_hash` and invalidate every gate ledger, so each run dir resumed afterwards needs a golden
    re-run (4.5). Accept one re-run at M7/M8, or take the alternative that keeps `align/eval/` byte-identical?
12. **[A2] Chain manifest.** The freeze closure is only as complete as the list of entry points. Should the repo carry
    `tests/frozen_entrypoints.txt`, maintained by whoever writes a new `.tmp/p2/r1/chain*.sh`, or should the closure scan parse the
    chain scripts directly?
13. **[A2] `CodeEnv` binaries.** `prlimit` and whatever interpreter runs candidate code are external binaries; DEPENDENCY_POLICY
    requires a dated entry for those (as for Lean, 2026-10-12). Which languages should `CodeEnv` support first?