[Feat] Route FSDP weight all-gather through copy-engine symmetric memory

[wtr0504]

🗂️ PR Category

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

📝 Description

Add a copy-engine (CE) transport for SimpleFSDP weight all-gather, as an alternative to NCCL ring kernels.
Today enable_fullgraph_overlap hides weight gathers under upstream compute, but those gathers still occupy SMs and contend with the compute that is supposed to hide them. With fsdp_config.transport="copy_engine":

1. Materialize FSDP shards into a torch.distributed._symmetric_memory arena at model build (to_empty interception on the decorated class). One window per (dtype, process group) so dense FSDP and MoE edp meshes do not share a rendezvous.
2. Rewrite marked weight all_gather_into_tensor / _coalesced nodes whose input is an arena shard (to_local(placeholder|get_attr)) to magi::symm_all_gather / magi::symm_all_gather_coalesced. Cast/pad gathers stay on NCCL, so mixed graphs are legal.
3. Gather by peer copy-engine reads (cudaMemcpyBatchAsync) into a fresh dest — zero SM occupancy, no per-step cross-rank barrier. Weights are static in inference, so the NCCL barrier is wasted work.
4. Reuse the existing latest-safe-launch reorder: CE ops lower to FallbackKernel, so the pass now treats magi::symm_all_gather* as collectives (_issues_transfer) instead of counting them as compute.
5. Cost-model the gather as wait(launch()). Timing the launch alone measures CPU issue (~3 µs) and misses the side-stream copies, which undersizes the overlap window by an order of magnitude.
   Default remains transport="nccl". CE requires NVLink within the FSDP mesh dim and static weights (inference).

Config

Option	Default	Purpose
fsdp_config.transport	"nccl"	"nccl": SM ring kernels. "copy_engine": symmetric-memory peer copies.
Enable with existing fullgraph overlap:

compile_config.disable_graph_split = True
compile_config.cudagraph_mode = CudaGraphMode.NONE
compile_config.fsdp_config.enable_fsdp = True
compile_config.fsdp_config.transport = "copy_engine"

[jiahy0825]

Rewrite this file as a custom pass like magi_compiler/passes/full_graph/replace_sage_atten.py

[jiahy0825]

The word arena is weird in this pr, you can use SymmBuffer directly. Replace all arena with a better name

[jiahy0825]

It makes no sense to apply a patch when we use ce-ag. I think it's better to replace all weights with SymmBuffer after capturing the whole graph.

[jiahy0825]

Another review comment from ai~