Price CP2 recompute per rank on each layer type's layout

[bradhilton]

Follow-up to #963 (stacked on it). At CP2, TrainerRank's checkpoint floor now prices each context-parallel rank on its own layouts, instead of pricing every rank on the busiest rank's rows with one mixer width. This makes the estimate correct per rank rather than tighter: on the measured workloads its accuracy is about unchanged.

Why: the two CP2 ranks don't hold the same work.

• On a real-data batch, rank 0 held 52,480 attention rows and rank 1 44,314. Rank 1's attention still kept more memory for backward (5.76 GB vs 4.15 GB), because its remote stage covers all its queries, so it runs two full-length stages.
• The old floor charged both ranks 5.00 GB for attention. That was 0.76 GB under for rank 1, hidden by over-charges on boundaries and GDN.

What changes (CP2, TP1, ART's CP core attention without a softmax offset, gradient groups; everything else keeps the old floor):

• Saved boundaries use the layout each layer's input arrives in. GDN layers that follow another GDN layer take GDN-layout rows; every other layer takes attention-layout rows.
• A recomputed attention layer is priced on attention rows, plus what the CP executor keeps for backward. That comes from a size-only mirror of the executor (retained_stage_record_bytes), kept next to the executor code.
• A recomputed GDN layer is priced on GDN rows.
• The MoE stage uses each layer's own rows for its local terms. Routed rows are the EP share when the EP group is the CP group, and otherwise the busiest rank's rows, as before. The MoE workspace no longer clamps routed rows to a rank's local rows, since a rank can receive more than it holds.
• Plan pricing and admission take the largest rank. Building the peer rank's CP plan costs about 40–76 ms the first time a batch layout is planned (under 1 ms once cached), now counted in planning time. (Under CP, width search already admits on materialized plans.) The split search's lower bound prices even-share layouts at the least attention state, so it stays below the per-rank cost.

Measured against the PyTorch peak on the most-loaded rank, EP2/CP2 on two H200s. Peaks and losses are identical before and after; only the estimate changes. A later fix counts one more LSE buffer per attention stage, which adds under 0.01 GB, inside this table's rounding.

case	before (cold / warm)	after (cold / warm)
real data, 40 layers	+18.2% / +19.2%	+18.2% / +20.2%
random weights, 8 layers	+26.5% / +23.6%	+26.6% / +23.7%
random weights, 40 layers	+19.7% / +17.5%	+18.8% / +16.6%

The real-data warm call rises because its rank 1 again runs two full-query stages, which the old floor under-priced.

Limitation. The executor mirror and the layout ledger are validated at CP2 against allocator traces of Qwen3.6-35B-A3B, within 0.3%. No test yet runs the executor itself against the mirror; that needs GPUs. CP4 and other executor configurations keep the old floor. GDN keeps its existing per-row width, now on GDN rows. The EP routing allowance, which accounts for most of the remaining EP2 gap, is unchanged.

Testing: trainer-rank and CP unit tests pass; the table's runs are on local H200s.

🤖 Generated with Claude Code