[Feat] Keep FSDP overlap when per-rank graphs diverge in compute

[wtr0504]

🗂️ PR Category

• ✨ New Feature
• 🚀 Optimization (performance, memory, etc.)
• 💥 Breaking Change
• 🐛 Bug Fix
• 🛠️ Development / Refactoring
• 📚 Documentation
• 🧹 Chore (Dependencies, CI/CD, Configuration, etc.)
• 🧪 Testing

📝 Description

FSDP fullgraph overlap no longer requires per-rank graphs (or profiler key sets) to be identical. Ranks may diverge in pure compute while the NCCL-issuing sequence (weight AG / CP / EP) stays the same. This PR keeps overlap on in that case, instead of aborting the reorder or falling the whole cost table back to analytical.

Reorder: negotiate placement relative to the collective skeleton

• Build a rank-comparable skeleton of every snode that issues NCCL (snode_issues_collective: functional collectives + custom ops flagged has_internal_collective).
• _negotiate_mode picks one of identical / slot / pinned / abort from a symmetric all_gather:
  • slot: graphs differ, skeleton matches → place each gather in a consensus skeleton slot (may hop CP/EP if every rank hops the same one); in-slot index stays per-rank.
  • pinned: skeleton itself differs → pin each AG between the NCCL-issuing snodes it already sat between (no hop).
  • abort: weight-AG counts differ → leave the graph unchanged.
• Slot targets are max-reduced and clamped non-decreasing so gathers cannot swap. Commit of the rebuilt order is AND-reduced across ranks.

Profiler: lockstep-measure the key intersection

• warm_and_sync measures keys present on every rank that also have a stashed snode (same sorted order, fixed iters), then max-reduces.
• Rank-local pure compute keeps its warmup measurement; rank-local collective-bearing keys fall back to analytical (solo NCCL replay would hang).
• Extern replay supports float8 (cast from fp32 noise).

Custom-op materialize-input API (profiling/materialize_inputs.py)

Replay of a custom op first generic-realizes every argument, then optionally runs a same-signature hook to rebuild value-dependent metadata (cp_split_sizes, cu_seqlens, ...). The hook is no longer (fx_node, realize) -> (args, kwargs):

def _attn_cp_inputs(q, k, v, cp_split_sizes):
    seq = int(q.shape[1] if q.dim() == 4 else q.shape[0])
    return q, k, v, [seq] * len(cp_split_sizes)

@magi_register_custom_op("mylib::attn_cp", materialize_inputs=_attn_cp_inputs, has_internal_collective=True)
def attn_cp(q, k, v, cp_split_sizes): ...

• register_materialize_inputs(op_name, fn=None, *, has_internal_collective=False) — fn is optional when only the collective flag is needed.
• apply_materialize_inputs(hook, args, kwargs) — runs the hook on already-realized args; None keeps the generic realize.
• get_materialize_inputs_hook / op_has_internal_collective — lookup.
• magi_register_custom_op(..., has_internal_collective=, materialize_inputs=) wires both at registration time.

Tests

• Reorder mode-ladder self-check (identical / slot / pinned / abort) and world=2 mismatch now asserts SLOT + identical final collective sequence.
• Estimator mismatch path asserts shared keys stay measured.
• Registry / custom-op decorator coverage for same-signature materialize_inputs hooks and flag-only has_internal_collective.