[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. Bind eligible FSDP weights into torch.distributed._symmetric_memory windows, driven from the captured FX graph between lowering and bucketing, using the live tensors Dynamo handed the backend. Windows are pooled per (process group, dtype) and capped at 4 GiB, so dense FSDP and MoE edp meshes never share a rendezvous. The publish barrier is paid once at bind, never per step.
2. Gate on what the copy engine can actually serve: a single even Shard(0) DTensor, contiguous, on CUDA, whose gather input is exactly to_local(placeholder|get_attr). An uneven Shard(0) is refused — the gather copies one fixed-size slab per peer, so a rank whose peers own fewer rows would read past the end of theirs. The decision is agreed across ranks before anything moves: a gather retargeted on only some ranks would never complete.
3. Rewrite the marked gathers to magi::ce_all_gather / magi::ce_all_gather_coalesced. Bucketing splits on the bind flag so a bucket is all-bound or all-unbound; cast/pad gathers and anything left unbound stay on NCCL, so mixed graphs are legal.
4. Gather by peer copy-engine reads (cudaMemcpyBatchAsync, one submission per bucket) 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.
5. Reuse the existing latest-safe-launch reorder: CE ops lower to FallbackKernel, so the pass now treats magi::ce_all_gather* as collectives (_issues_transfer) instead of counting them as compute.
6. 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.
fsdp_config.symm_min_shard_mib	0	Size floor for binding. Local shards below it stay on NCCL; 0 binds every eligible weight.
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]

Another review comment from ai~

[jiahy0825]

LGTM