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

.github/workflows/prek.yml

@@ -241,6 +241,8 @@ jobs:
             tests/unit/test_trainer_rank_pending_memory.py \
             tests/unit/test_trainer_rank_shared_memory.py \
             tests/unit/test_trainer_rank_converted_memory.py \
+            tests/unit/test_trainer_rank_layout_memory.py \
+            tests/unit/test_context_parallel_retained_bytes.py \
             tests/unit/test_trainer_rank_split.py \
             tests/unit/test_megatron_compile_garbage.py \
             tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \
@@ -287,4 +289,6 @@ jobs:
             --ignore=tests/unit/test_trainer_rank_pending_memory.py \
             --ignore=tests/unit/test_trainer_rank_shared_memory.py \
             --ignore=tests/unit/test_megatron_compile_garbage.py \
-            --ignore=tests/unit/test_trainer_rank_converted_memory.py
+            --ignore=tests/unit/test_trainer_rank_converted_memory.py \
+            --ignore=tests/unit/test_trainer_rank_layout_memory.py \
+            --ignore=tests/unit/test_context_parallel_retained_bytes.py

src/art/megatron/context_parallel/executor.py

@@ -35,6 +35,7 @@
     DkvReducePlan,
     ExactMaskMetadata,
     FlexMaskSpec,
+    RankRuntimePlan,
     StageExecutionSpec,
     StagePlan,
     TokenRange,
@@ -1550,6 +1551,112 @@ def _merge_stage_output_grads_from_tape(
     return stage_out_grads, stage_lse_grads
 
 
+def minimum_retained_bytes_per_row(
+    *,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    value_head_dim: int,
+    element_size: int,
+) -> int:
+    """The least ``retained_stage_record_bytes`` keeps per own row.
+
+    Every rank with rows runs a local stage over all of them (each row attends
+    to itself); an aligned one keeps at least the contiguous copies of
+    multi-head views, flex's output and its two LSEs.
+    """
+    copies = (q_heads > 1) * q_heads * head_dim + (kv_heads > 1) * kv_heads * (
+        head_dim + value_head_dim
+    )
+    return (copies + q_heads * value_head_dim) * element_size + 2 * q_heads * 4
+
+
+def retained_stage_record_bytes(
+    rank_plan: RankRuntimePlan,
+    *,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    value_head_dim: int,
+    element_size: int,
+    block_size: SparseBlockSize,
+) -> int:
+    """Bytes a recomputed attention keeps for backward beyond its own-row tensors.
+
+    A size-only mirror of ``_forward_stage_records`` with ``record_for_backward``
+    and of ``_run_stage_attention``; keep the three in step. Every stage that
+    runs keeps the Q/K/V flex consumed: a padded copy when the execution length
+    differs, else a contiguous copy of the permuted ``q_flat``/``k_flat`` view;
+    partial-range gathers and remote fetch buffers are kept as the stage's
+    inputs. It keeps flex's output and LSE at the execution length, and
+    logical-length copies of them when padded, plus flex's own LSE beside the
+    normalized one the FLASH backend returns (counted on every backend). Every
+    producing stage after the first keeps a merge-tape clone of the
+    accumulators. Accumulators themselves are transient.
+    """
+    own = int(rank_plan.local_valid_lengths[0]) if rank_plan.local_valid_lengths else 0
+    accum_size = 4 if element_size < 4 else element_size
+    q_row = q_heads * head_dim * element_size
+    k_row = kv_heads * head_dim * element_size
+    v_row = kv_heads * value_head_dim * element_size
+    out_row = q_heads * value_head_dim * element_size
+    lse_row = q_heads * 4
+    tape_row = q_heads * (value_head_dim + 1) * accum_size
+    total = 0
+    local_produced = False
+    tapes: list[int] = []
+    for stage in _ordered_stage_plans(rank_plan.stage_plans):
+        if not (stage.q_len > 0 and stage.k_len > 0 and stage.slices):
+            continue
+        q_len = _logical_stage_q_len(stage)
+        k_len = _logical_stage_k_len(stage)
+        q_pad, k_pad, _family = select_sparse_execution_family(
+            is_local_stage=bool(stage.is_local_stage),
+            q_len=int(stage.q_len),
+            k_len=int(stage.k_len),
+            block_size=block_size,
+        )
+        q_full = _ranges_cover_full_length(stage.owner_local_q_ranges, length=own)
+        # Queries: the full range is a permuted view of q_flat; a partial range
+        # is gathered into a contiguous tensor kept as the stage's input.
+        if not q_full:
+            total += q_row * q_len
+        if q_pad != q_len:
+            total += q_row * q_pad
+        elif q_full:
+            # A view of the projection output: flex copies it unless it is
+            # already contiguous, which one head alone does not guarantee
+            # (a fused QKV split keeps the projection's token stride).
+            total += q_row * q_len
+        # Keys and values: local ranges as for queries; remote ones land in
+        # contiguous head-major fetch buffers kept as the stage's inputs.
+        k_full = bool(stage.is_local_stage) and _ranges_cover_full_length(
+            stage.owner_local_k_ranges, length=own
+        )
+        if not k_full:
+            total += (k_row + v_row) * k_len
+        if k_pad != k_len:
+            total += (k_row + v_row) * k_pad
+        elif k_full:
+            total += (k_row + v_row) * k_len
+        total += (out_row + 2 * lse_row) * q_pad
+        if q_pad != q_len:
+            total += (out_row + lse_row) * q_len
+        tape = tape_row * (own if q_full else q_len)
+        if not q_full:
+            # Its int64 row index is kept even by the first producing stage.
+            total += 8 * q_len
+        if stage.is_local_stage:
+            local_produced = True
+        else:
+            tapes.append(tape)
+    # The first producing stage records no tape: the local stage when it ran,
+    # else whichever remote stage is ready first, so drop the smallest tape.
+    if tapes and not local_produced:
+        tapes.remove(min(tapes))
+    return total + sum(tapes)
+
+
 def _forward_stage_records(
     *,
     q_flat: torch.Tensor,

src/art/megatron/context_parallel/runtime.py

@@ -433,6 +433,55 @@ def context_parallel_rank_model_token_counts(
     )
 
 
+def context_parallel_rank_layouts(
+    *,
+    group_ids: torch.Tensor,
+    parent_ids: torch.Tensor,
+    topology: ParallelTopology,
+    config: ContextParallelConfig,
+    original_seq_len: int,
+    build_gdn_execution_spec: bool,
+    gdn_planner_config: Any | None = None,
+) -> tuple[tuple[int, ...], tuple[int, ...] | None, tuple[RankRuntimePlan, ...]]:
+    """Each CP rank's attention rows, GDN rows and attention stage plan.
+
+    Uses the cached planning bundle and per-rank runtime plans that execution
+    builds, so a memory estimate sees the layouts the ranks will run.
+    """
+    planning_key, bundle, _group_ids_cpu, _parent_ids_cpu = (
+        _get_or_build_planning_bundle(
+            group_ids=group_ids,
+            parent_ids=parent_ids,
+            topology=topology,
+            config=config,
+            original_seq_len=original_seq_len,
+            build_gdn_execution_spec=build_gdn_execution_spec,
+        )
+    )
+    attention = tuple(bundle.token_layout_index.token_counts_by_rank)
+    gdn = None
+    if build_gdn_execution_spec:
+        gdn = tuple(
+            _plan_gdn_global_execution(
+                planning_key=planning_key,
+                bundle=bundle,
+                topology=topology,
+                gdn_planner_config=gdn_planner_config,
+            ).gdn_token_counts_by_rank
+        )
+    plans = tuple(
+        _get_or_build_bundle_rank_plan(
+            planning_key=planning_key,
+            bundle=bundle,
+            original_seq_len=original_seq_len,
+            target_rank=rank,
+            block_size=config.block_size,
+        )
+        for rank in range(len(attention))
+    )
+    return attention, gdn, plans
+
+
 def context_parallel_model_token_total(
     *,
     group_ids: torch.Tensor,