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

[bradhilton]

With #1002 (merged in through #963), each CP rank also prices #1002's adapter-gradient term on its own layouts:

• Per-rank boundaries. Each rank walks its own saved boundaries, one gradient group after another.
• Paired floors. The extra is paired with that rank's own floor.
• Why paired. A single short CP2 sequence can leave one rank with no GDN rows. That rank's extra is then almost every adapter gradient, and adding it to the other rank's floor over-counted by up to 632 MB.
• Result. On the full stack (Qwen3.6-35B-A3B, 40 layers, CP2, one real sequence per wave), cold first waves come out +0.9% to +2.0% raw at EP1 and +3.0% to +13.4% at EP2. None are under.

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.

Planner-miss reports replay the per-rank pricing on CPU. Each group's CP layouts and the decoder's per-layer input layouts are recorded with the selected plan.

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

🤖 Generated with Claude Code