perf(glu): emit residual taps from the clean-forward scan; one frozen forward per step

[claude-spd1]

Description

Bridge task: scan-read-activations-taps (2026-07-10 JAX compile-hygiene audit finding).

read_activations was a Python for layer in range(n_layer) — the largest unrolled region left in the step HLO, inlined into the train step and all seven jitted eval/metric steps — and a duplicate of the scanned clean_output forward on the same batch two lines earlier in the train step (XLA won't CSE scan vs unroll). Authored against llama8b.py; after #987 the machinery lives in targets/glu_transformer.py, so the fix now covers llama8b and qwen3 at once (rebased onto a2c297b).

Now clean_output / read_activations / the new clean_output_and_activations all route through one private _clean_forward lax.scan (SPEC S3/S4 semantics unchanged):

• Taps ride out as scan ys — one stacked [tapped_depth, b, t, d] array per block intermediate (resid / post-LN1 / attn-out / post-LN2 / MLP-inner), indexed into the flat tap dict after the scan (the _run_masked_forward per-kind-ys pattern). No per-layer lax.cond (the scan-recon-chunks lesson). Family behavior (Qwen3 QK-norm) rides in attn.core exactly as the unrolled version had it.
• Only the block prefix up to the last tapped layer emits ys; the remainder scans the same body emitting nothing — the old early-break compute property, preserved statically.
• clean_output_and_activations (new DecomposedModel protocol method; trivial two-call impl on SimpleMLP/TMS/ResidMLP) serves the two steps that need clean logits AND taps on the same batch — train and fast eval — from ONE frozen forward. Harvest's load_run forward folds its three frozen passes (clean + CI taps + site inputs) into one fused call likewise.

Measured (32L/224-site structure at toy width, canonical 4-chunk recon, both arms one CPU node): HLO text 9.91→9.49 MB (−4%), CPU compile 54.2→49.9s. Per the 2026-06-30 review doc the step-time win is small (~0.3%: read_activations gathers no V/U and the step is gather-bound) — the payoff is HLO/compile across all eight jitted steps plus one whole frozen forward removed per train/eval step. A 4-node dp32 GPU mem_profile A/B (compile time + peak memory, per the task brief's memory caveat) is in flight; numbers land on the bridge thread and here.

Related Issue

Bridge thread scan-read-activations-taps. Coordinated with targets-refactor (#966): agreed there that this PR's scan shape is what their remaining carve builds on.

Motivation and Context

Compile hygiene: everything else on the layer axis is already lax.scan'd (clean_output, masked forwards, chunkwise CI fn). This removes the last per-layer unroll and a redundant frozen forward per jitted step.

How Has This Been Tested?

• Pinned pre-refactor CPU fixtures (557 arrays: clean logits, taps for 4 wanted shapes, 3-step train trajectories in BOTH remat modes, eval-step metrics), re-pinned at the rebased base a2c297b: clean_output logits bit-identical; taps vs the old unrolled forward within fp32 reassociation (≤1.7e-5 abs — the scan-vs-unroll class clean_output's docstring already documents; same order as SPEC D4's accepted 2.2e-6), propagating to ≤3.9e-6 abs over the 3-step trajectory. Taps are now bit-identical to the residual the recon target actually threads (previously they came from a separately-compiled forward).
• New test test_clean_output_and_activations_shares_the_forward: fused accessor EXACTLY equals the separate methods (and taps are invariant to the wanted set); at full tapped depth (every production config) fused clean logits are bit-equal to clean_output.
• Full suite on a SLURM CPU node: pytest green at default devices and --xla_force_host_platform_device_count=4, basedpyright 0 errors, ruff clean (re-running post-rebase; will confirm on the thread).

Does this PR introduce a breaking change?

DecomposedModel gains one protocol method (clean_output_and_activations); all in-repo targets + test stubs implement it. read_activations now asserts wanted is non-empty (previously crashed on max() of empty). No config or checkpoint format changes.

🤖 Generated with Claude Code

Crew-Address: task/scan-read-activations-taps

[claude-spd1]

GPU memory A/B done — the task brief's remaining verification (taps-as-scan-ys are saved outputs; confirm peak memory).

Setup: two mem_profile arms on the canonical dp32 HSDP topology (4 nodes × 8×H100-80GB), llama8b_full32L_HSDP_b32_dp32 shape, identical config/hardware; before = base a2c297b (run p-7f38afc9, job 1090542), after = this PR's 0529fb8 (p-2146aea8, job 1090543).

XLA static memory_analysis() of jit_step:

	argument	output	temp	alias
before	28.12 GiB	24.54 GiB	47.33 GiB	24.54 GiB
after	28.12 GiB	24.54 GiB	47.23 GiB	24.54 GiB

No regression — temp down 0.10 GiB. As the brief anticipated, the [n_layer,B,T,d] residual stacks were already resident as step outputs under the unrolled version; emitting them as scan ys changes nothing there, and removing the duplicate frozen forward trims transients slightly.

(Runtime 2-step peaks weren't obtainable on this hardware: argument+temp = 75.4 GiB exceeds the 73.6 GiB BFC pool at the default 0.92 mem fraction, so BOTH arms OOM identically at step 1 — b32/dp32 is the minimum per-rank batch, and the b128 config of the 2026-06-29 memory canon is further out of reach on 80GB parts. The static analysis is XLA's own accounting of the same executables and is the apples-to-apples number.)

Together with the earlier CPU A/B (HLO text −4%, compile 55.5→49.7s at the canonical 4-chunk 32L structure), all verification from the task brief is complete.

Crew-Address: task/scan-read-activations-taps