Apply trainer-rank memory floors per context-parallel rank

[bradhilton]

Partially addresses #949.

Outside TP1/CP1/EP1, every model-aware memory floor was disabled:

• _moe_output_bytes_per_token returned 0 unless the shape was TP1/CP1;
• _checkpoint_memory_floor returned (0, 0) unless the topology was (1, 1, 1) with EP (1, 1);
• _head_workspace_bytes was gated the same way.

So with no profile yet, a CP forward fell back to the generic activation factor, packed × hidden × 2 B × activation_factor. In the #949 CP2/EP2 report, the admitted 44.68 GB prediction is exactly that fallback: 779,067 × 2,048 × 2 × 14.

Measurements

These are component measurements: local H200s, Qwen3.6-35B-A3B, random weights, EP1, forward_micro_batches with forward and backward over ten unique-token requests scaled from #949's first wave (about 195K tokens). They are not full-depth HybridEP qualification.

layers	CP	observed cold peak per rank	per local token	main prediction	this PR
8	1	95.3 GB	~490 KB	69.5 GB	69.5 GB
8	2	46.0–51.9 GB	472–532 KB	5.27 GB	29.4 GB
8	4	23.3–26.8 GB	478–550 KB	5.27 GB	15.0 GB
40	4	29.4–33.5 GB	603–688 KB	~11.2 GB (generic)	30.4 GB

Memory per local token is essentially the same at CP1, CP2 and CP4 (about 470–550 KB cold), so the CP1 floors apply per rank. This PR raises cold CP predictions 3–6×, but they are not yet safe cold bounds:

• At 8 layers, predictions are 0.56–0.64× of observed at CP2/CP4, against 0.73× at CP1.
• Part of the gap is the GDN pending floor (_gdn_memory.model_shapes / plan_floor), which is still CP1-only. CP predictions per local token (~301 KB) are therefore below CP1's (~356 KB), and warm second calls still exceed them by up to about 22% (CP2 36 vs 29.4 GB; CP4 18.5 vs 15.0 GB).
• The rest is a cold first-call excess that also exists at CP1 (69.5 vs 95.3 GB) and scales with tokens. It is consistent with compile/autotune workspaces on these long sequences, but that attribution is a hypothesis.
• At 40 layers on CP4, the prediction is 30.4 GB against 29.4–33.5 GB observed. The most loaded rank exceeds it by 10.2%. Since the prediction already includes the 1.1× safety factor, that rank's pre-safety estimate is about 21% low.

Changes

• _plan_group_rows prices each group at the rows of the most loaded CP rank (_max_rank_model_tokens, the same assignment _plan_retained_tokens already uses). The split lower bound uses an even share, ceil(rows / cp), which never exceeds that maximum.
• _checkpoint_memory_floor and _head_workspace_bytes require only TP1/PP1. Checkpoint retention also applies at EP > 1.
• _moe_output_bytes_per_token accepts CP > 1, since CP shards rows rather than the per-token working set. TP, EP and ETP still return 0.
• At CP > 1 the static MoE workspace term uses per-rank group rows instead of global packed tokens.
• Under CP, width probes defer to the existing exact-plan fallback, because global packed counts would price the per-rank floors about CP× too high. Without this, a probe could reject widths the exact planner admits.
• The exact-plan fallback retries the memory-minimal (full-sharing) layout both when the cost-optimal layout does not fit and when it falls outside the profile's trust window, as main's cheap probe did. The minimum wave gets the same trust-window retry. Every retry condition is reduced across DP ranks, so all ranks take the same branch. Under CP, full sharing can price higher than the cost-optimal layout, because the busiest rank's rows are not monotone in packed tokens and full sharing has more GDN segments. If the retried layout does not fit, the minimum wave goes to the split ladder, which re-plans the cost-optimal unsplit layout first and admits it. The result matches main; the cost is two extra plan-and-price rounds.
• Planning cost: under CP, every priced plan runs the CP rank assignment per group. That includes rejected width probes, split-ladder rungs and planner snapshots, and each deferred probe also runs the full planner. main skipped this during admission for CP with full recompute, MoE or no_grad. A memory check computes the assignment twice per group; the second is a planning-bundle cache hit that still re-digests the metadata.

Not addressed here

• EP > 1 routed-expert working set: the HybridEP "flex" dispatcher is not recognized, so it is still unmodeled. TrainerRank CP2/EP1 normal admission reaches CUDA OOM in compiled MoE LoRA forward #949's EP2/EP4 runs therefore gain the checkpoint retention floors but not the MoE term.
• GDN pending floor under CP: _gdn_memory.plan_floor stays CP1-only; see the shortfall above.
• The cold first-call excess: it exists at CP1 too.
• Evidence limits: these are random-weight EP1 component runs. They do not cover pretrained routing, HybridEP, adversarial sharing, split admission, TP/PP > 1 or DP > 1.

Validation

• CI on the current head: quality-checks and the 2xH200 TrainerRank GPU validation pass. Earlier GPU runs failed at launch for lack of 2xH200 capacity (ResourcesUnavailableError), before any test ran.
• CPU-only (CUDA_VISIBLE_DEVICES=""), tests/unit/test_trainer_rank_*.py plus test_prefix_tree_packing.py: 1,242 passed on the reviewed code. The planning-status fixture now gives its stub signature a CP1 topology, because group rows read the plan's CP size; that file passes 9/9. The other failures are timing-sensitive distributed tests that pass in isolation on this head and on main, or flake on both.
• Updated tests: the tests that encoded the CP exclusion now assert CP support, while TP/PP (and EP for MoE) still disable the floors.
• New regressions (TP1/CP2):
  • With even and uneven rank ownership, the width probe defers, the per-rank checkpoint floor is active, and the split lower bound stays at or below exact required and retained costs.
  • Width search reaches width 15 by retrying full sharing inside the trust window; it stops at 8 without the retry. A minimum wave outside the trust window is admitted warm with full sharing.
• Lint and types: ruff and ty are clean on the changed files.

🤖 Generated with Claude Code