perf(recon): prefix-shared masked forwards — one clean pass seeds every recon forward

[danbraunai-goodfire]

Description

Measured speedup (dp8 single-block smokes, B200, 300-step A/B at identical config + seed):

run	before	after	speedup	peak HBM
layer-18 block (all 7 matrices, doubled C)	1.477 s/step	1.18 s/step	1.25×	49.0 → 44.6 GB
layer-31 block	1.186 s/step	0.652 s/step	1.82×	43.7 → 42.1 GB

The win scales with block depth (fleet average over 32 per-block runs ≈ 1.20×), and chunkwise multi-block runs get each chunk pass truncated at its own chunk start. Losses match the baseline runs to 4 significant figures in early steps; the layer-31 trajectory is identical through all 300 steps, layer-18 within 0.5% at step 300 (reassociation amplification, D5).

Where the speedup comes from: a step previously ran the frozen prefix of the target ~6× — once in clean_output, once in read_activations, and once inside each masked forward (2 PPGD warmup ascents + stochastic recon + PPGD main), which all re-embedded the batch and re-ran every block below the first live site. Now the step runs ONE clean pass (clean_output_and_taps, a segmented scan) that emits the recon target, the CI taps, and the residual boundaries; every masked/ascent forward starts from its entry's boundary (masked_start → (first_live_block, resid)), so the frozen prefix below the earliest live site is never recomputed. For a single-block-18 decomposition that deletes 4–5 × 18 frozen-block forwards per step.

Mechanically: the DecomposedModel masked family (masked_output / masked_output_stochastic / masked_site_outputs / masked_component_activations) takes a model-opaque forward start instead of the token batch, built via two constructors — start_from_inputs (input edge; eval tiers) and masked_start (tap-seeded; all recon/ascent forwards). The engine stays d-blind: starts are built and consumed target-side (the S4 opacity rule), start_block is static per recon-plan entry, and all-sites passes on a full-model run degrade gracefully to start_block = 0. All seven DecomposedModel impls migrated (toys: identity — their start IS their input). llama8b's already-segmented [prefix][live][suffix] scan just gets its prefix sliced [start_block, first_live); llama_simple_mlp destructures the start's static int OUTSIDE its jax.checkpoint (tracer hazard, pinned by test).

This seam is also the substrate for later activation caching (resid from disk instead of the clean pass) and adjacent to block-local recon experiments.

Related Issue

None — motivated by the per-block (two-stage) training investigation; step-cost model measured in wandb group l18only (runs 172591/172594-172596, A/B 172602/172603).

Motivation and Context

Per-block VPD runs (the l18-style runs and the planned per-block stage-1 sweep) and mid-stack chunkwise runs (l18-26) spend a large share of each step recomputing a frozen prefix whose value the clean forward already produced.

SPEC (needs @oli sign-off per the one rule): S18 amended — masked forwards take a forward start; the clean path still takes the token batch and embeds internally, clean_output stays the whole-model frozen forward (S3 untouched), and NONE of the removed residual-start engine machinery returns (no Prefix object, no first_layer offsets, no out-of-graph harvest, no tap ceiling — cf. the 2026-06-24 removal). New D5: tap-start ≡ input-start up to float reassociation. S2/S3/S10′/S24 untouched. Glossary + §4.1/§4.5 pseudocode updated.

One drive-by fix: experiments/invariance_check.py had been bitrotted since the 2026-06-24 token contract (fed a float residual batch into embed[tokens]) — now feeds tokens and passes.

How Has This Been Tested?

• New tests/test_prefix_start.py (the D5 gate): tap-start vs input-start equivalence on the deterministic AND stochastic mask paths (stochastic RNG keys fold over absolute layer indices, so draws are start-invariant), both LM targets, remat on/off; clean_output_and_taps vs separate calls (clean bit-level, taps at D4 tolerance — scan vs unrolled loop).
• Full param_decomp/tests/ + toy suites: 336 passed at default AND --xla_force_host_platform_device_count=4 (skips/xfails pre-existing).
• experiments/invariance_check.py (D4): 3-step trajectory matches 1-layout vs 4-device GSPMD, reassociation-only.
• make type / make check clean.
• Live cluster A/B above (branch perf/prefix-shared-recon-ab = this + the smoke configs); losses vs baseline as described.

Does this PR introduce a breaking change?

Yes, two contract-level notes:

• The masked-family signature changed everywhere (no legacy path): callers pass a forward start built by start_from_inputs / masked_start. All in-repo call sites are migrated; out-of-tree consumers of masked_output must wrap their batch in model.start_from_inputs(batch).
• A mid-run resume across this change is trajectory-continuous only up to float reassociation (CI taps moved from the unrolled read_activations pass to the clean scan's boundary carries), and the persistent compile cache misses once (new HLO).

🤖 Generated with Claude Code

https://claude.ai/code/session_011F8uqqWqnkC48Y8G4QCKsk

[claude-spd1]

Crew review (task review-pr-965-prefix-shared-recon). Recommend merge once Oli signs off the S18/D5 amendment — that gate is the only blocker I see.

Verified:

• Read the full diff; re-ran the D5 gate (test_prefix_start.py) plus test_llama8b / test_llama_simple_mlp / test_eval / test_generic_model_io / test_attn_patterns_eval / test_hidden_acts_eval / test_source_grad_mean at cfc6f65c9 locally: 53 passed. CI green.
• The seam respects the 2026-06-24 removal's substance: starts are built and consumed target-side, the engine never inspects one (S4-style opacity), no Prefix/first_layer/harvest machinery returns. clean_output stays the whole-model frozen forward (S3), and wanted=() makes clean_output_and_taps literally the old unsplit scan.
• _run_masked_forward's start_block <= first_live assert + the contiguous-live assert cover the seeding contract; live=() correctly degrades to a frozen scan from start_block.
• Stochastic start-invariance rests on absolute-layer key folding and is pinned by test.

Findings (all minor, none blocking):

1. targets/llama8b.py class docstring (~line 512) names start_from_taps — the method is masked_start.
2. eval.py builds masked_start(token_ids, taps, site_names) from taps requested as bare ci_fn.input_names, relying implicitly on the chunkwise builder's guarantee that the first chunk taps resid.{earliest decomposed block}. True today (_resolved_chunks), and a violation fails loud (KeyError at trace), but train.py unions start_taps into wanted — eval.py doing the same one-line union would make the invariant structural instead of coincidental.
3. Config-space narrowing: an LM target + layerwise_mlp/global_mlp CI fn (site-name input_names) was previously traceable and now trips the resid.*-only assert in clean_output_and_taps at trace time. The config union permits the combo though the docstrings scope MLP CI fns to positionless toys — if it's meant to be illegal, a guard at config resolution would say so sooner and clearer.
4. The chunkwise multi-block claim ("each chunk pass truncated at its own chunk start") is covered by tests but not by the live A/B (both A/B runs are single-block). Worth one chunkwise smoke before leaning on the number for chunk-run capacity planning.

On the two flagged veto-able choices, my recommendation is to keep the PR's current shape on both:

• Keep the clean_output one-line delegator. clean_output is S3's name for a real concept; clean_output_and_taps(x, ())[0] at 15 call sites trades a one-liner for noise everywhere.
• Keep start_from_inputs + masked_start separate. They have different signatures, different call sites (eval tiers vs recon/ascent), and collapsing them means an optional-taps mega-constructor — worse, not simpler.

Crew-Address: task/review-pr-965-prefix-shared-recon