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

magi_compiler/config.py

@@ -256,14 +256,29 @@ class FSDPConfig(BaseModel):
         ),
     )
     comm_overlap_window_scale: float = Field(
-        1.5,
+        1.0,
         ge=1.0,
         description=(
             "Multiplier on each collective's estimated runtime when sizing its compute window "
             "(need = comm * scale + margin): collectives are measured in isolation but run concurrent "
             "with the compute that hides them (~1.4-1.5x slower in-situ on 8xH100)."
         ),
     )
+    transport: Literal["nccl", "copy_engine"] = Field(
+        "nccl",
+        description=(
+            "Weight all-gather path. 'nccl': SM kernels. 'copy_engine': bind weights into "
+            "symmetric memory and gather with peer copy-engine reads (0 SM, no per-step barrier). "
+            "Needs NVLink on the FSDP mesh dim and static weights; unbound weights stay on NCCL."
+        ),
+    )
+    symm_min_shard_mib: int = Field(
+        0,
+        ge=0,
+        description=(
+            "Minimum local-shard MiB to bind for copy_engine. Smaller shards stay on NCCL. " "0 = bind every eligible weight."
+        ),
+    )
 
 
 def _find_cutlass_root() -> str:

magi_compiler/magi_backend/magi_backend.py

@@ -592,12 +592,13 @@ def _init_cache(self) -> str:
         self.local_magi_cache_path.mkdir(parents=True, exist_ok=True)
         self.compiler_manager.initialize_cache(self.local_magi_cache_path)
 
-    def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None:
+    def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule, example_inputs) -> None:
         """Whole-graph FSDP all-gather / compute overlap (disable_graph_split path).
 
-        1. Lower SimpleFSDP weight prim_redistribute -> explicit collectives and
-           optionally bucket them over the whole graph (no region partitioning;
-           buckets break only at dtype changes and the size cap).
+        1. Lower SimpleFSDP weight prim_redistribute -> explicit collectives, bind
+           the gathered weights into symmetric memory when the transport is the copy
+           engine, and optionally bucket them over the whole graph (no region
+           partitioning; buckets break only at dtype changes and the size cap).
         2. Install the profiling runtime estimator at ``config.estimate_op_runtime``
            (the analytical roofline is unusable for our sizing decisions).
         3. Install the latest-safe-launch reorder pass, REPLACING PyTorch's builtin
@@ -615,9 +616,17 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None:
         from magi_compiler.profiling import ProfilingRuntimeEstimator
 
         bucket_size_bytes = int(fsdp_cfg.bucket_size_mib) * 1024 * 1024
-        n_buckets = lower_and_bucket_full_graph(graph, fsdp_cfg.bucket_mode, bucket_size_bytes=bucket_size_bytes)
+        n_buckets = lower_and_bucket_full_graph(
+            graph,
+            fsdp_cfg.bucket_mode,
+            bucket_size_bytes=bucket_size_bytes,
+            transport=fsdp_cfg.transport,
+            example_inputs=example_inputs,
+            min_shard_bytes=int(fsdp_cfg.symm_min_shard_mib) * 1024 * 1024,
+        )
         magi_logger.info(
-            "FSDP fullgraph overlap: bucket_mode=%s bucket_size=%d MiB created %d buckets",
+            "FSDP fullgraph overlap: transport=%s bucket_mode=%s bucket_size=%d MiB created %d buckets",
+            fsdp_cfg.transport,
             fsdp_cfg.bucket_mode,
             fsdp_cfg.bucket_size_mib,
             n_buckets,
@@ -640,7 +649,7 @@ def _apply_fsdp_fullgraph_overlap(self, graph: fx.GraphModule) -> None:
         self.inductor_compile_config["reorder_for_compute_comm_overlap_passes"] = [reorder]
 
     @observe_lifecycle("graph_split")
-    def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[SplitItem]]:
+    def _split_graph(self, graph: fx.GraphModule, example_inputs) -> tuple[fx.GraphModule, list[SplitItem]]:
         # Step 1: resolve the splitting ops.
         if self.compile_config.disable_graph_split:
             assert (
@@ -658,7 +667,7 @@ def _split_graph(self, graph: fx.GraphModule) -> tuple[fx.GraphModule, list[Spli
 
         # Step 1.4: whole-graph FSDP overlap.
         if self.compile_config.fsdp_config.enable_fsdp:
-            self._apply_fsdp_fullgraph_overlap(graph)
+            self._apply_fsdp_fullgraph_overlap(graph, example_inputs)
 
         # Step 2: split graph by ops, we split graph based on resolved_ops, which becomes the partitioned single graph.
         subgraph_id = 0
@@ -719,7 +728,7 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> MagiSerializableFun
 
         self.full_graph_pass_manager(graph)
 
-        split_gm, piecewise_graphs = self._split_graph(graph)
+        split_gm, piecewise_graphs = self._split_graph(graph, example_inputs)
 
         submod_names_to_compile = [item.submod_name for item in piecewise_graphs if not item.is_splitting_graph]
         compilation_counter.num_piecewise_graphs_seen += len(piecewise_graphs)

magi_compiler/passes/fsdp_overlap/__init__.py

@@ -13,13 +13,24 @@
 # limitations under the License.
 
 from .bucket_all_gather import bucket_weight_all_gather_coalesced
+from .copy_engine import bind_weights_for_copy_engine, copy_engine_weight_candidates, rewrite_weight_ag_to_copy_engine
 from .lower_and_bucket import lower_and_bucket_full_graph
+from .node_meta import CE_BOUND, UNEVEN_SHARD, WEIGHT_AG, is_ce_bound, is_uneven_shard, is_weight_ag
 from .redistribute_lowering import lower_prim_redistribute_to_collectives
 from .reorder import FsdpOverlapReorder
 
 __all__ = [
+    "bind_weights_for_copy_engine",
     "bucket_weight_all_gather_coalesced",
+    "copy_engine_weight_candidates",
     "lower_prim_redistribute_to_collectives",
     "lower_and_bucket_full_graph",
+    "rewrite_weight_ag_to_copy_engine",
     "FsdpOverlapReorder",
+    "CE_BOUND",
+    "UNEVEN_SHARD",
+    "WEIGHT_AG",
+    "is_ce_bound",
+    "is_uneven_shard",
+    "is_weight_ag",
 ]

magi_compiler/passes/fsdp_overlap/bucket_all_gather.py

@@ -22,6 +22,8 @@
 
 from magi_compiler.utils import magi_logger
 
+from .node_meta import is_ce_bound, is_uneven_shard, is_weight_ag, mark_ce_bound, mark_weight_ag
+
 _ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default
 _ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default
 _WAIT = torch.ops._c10d_functional.wait_tensor.default
@@ -64,19 +66,18 @@ def _is_weight_all_gather(node: fx.Node) -> bool:
     """A SimpleFSDP weight ``all_gather_into_tensor`` launch."""
     if node.op != "call_function" or node.target is not _ALL_GATHER:
         return False
-    if node.meta.get("magi_fsdp_weight_ag"):
+    if is_weight_ag(node):
         return True
     return _gathers_a_weight(node)
 
 
 def _local_shard_bytes(ag: fx.Node) -> int:
-    """Bytes of the local (pre-gather) shard feeding a weight all_gather -- i.e. how
-    much this rank contributes to the collective.  Used to cap coalesced bucket size."""
-    loc = ag.args[0]
-    m = loc.meta.get("example_value")
+    """Bytes one rank contributes to a weight all_gather"""
+    m = ag.meta.get("example_value")
     if m is None:
         return 0
-    return int(m.numel()) * int(m.element_size())
+    _loc, world, _group = ag.args
+    return int(m.numel()) * int(m.element_size()) // int(world)
 
 
 def _split_by_dtype_and_size(
@@ -157,8 +158,9 @@ def _coalesce_one_bucket(graph: fx.GraphModule, node_index: dict[fx.Node, int],
     with graph.graph.inserting_before(first_ag):
         coalesced = graph.graph.call_function(_ALL_GATHER_COALESCED, (list(locals_), world, group_name))
         coalesced.meta["example_value"] = list(ag_metas)
-        coalesced.meta["magi_fsdp_weight_ag"] = True
-        coalesced.meta["magi_fsdp_weight_ag_coalesced"] = True
+        mark_weight_ag(coalesced, uneven=any(is_uneven_shard(ag) for ag in ag_nodes))
+        if all(is_ce_bound(ag) for ag in ag_nodes):
+            mark_ce_bound(coalesced)
 
         outs = []
         for i, am in enumerate(ag_metas):
@@ -177,7 +179,7 @@ def _coalesce_one_bucket(graph: fx.GraphModule, node_index: dict[fx.Node, int],
         graph.graph.erase_node(ag_old)
 
 
-def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: int = 0) -> int:
+def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes: int = 0, split_by=None) -> int:
     """Coalesce the SimpleFSDP weight all-gathers over the WHOLE graph: per process
     group, walk them in program order and cut a new bucket at every dtype change or
     when the accumulated local-shard bytes would exceed ``bucket_size_bytes``
@@ -193,22 +195,26 @@ def bucket_weight_all_gather_coalesced(graph: fx.GraphModule, bucket_size_bytes:
     are re-pointed from each old wait to ``wait_i``; the launch + getitems stay
     together so ``FsdpOverlapReorder`` later moves them as one unit.
 
+    ``split_by(node)`` adds a second bucket key, used in copy-engine mode to keep
+    bound and unbound gathers apart -- a bucket is one submission, so it cannot span
+    two transports.
+
     Runs after redistribute lowering (via ``lower_and_bucket_full_graph``).
     Returns the number of coalesced buckets created.
     """
     node_index = {n: i for i, n in enumerate(graph.graph.nodes)}
 
-    # Key by group_name only; dtype breaks buckets positionally inside
-    # _split_by_dtype_and_size (strict program-adjacency).
-    groups: dict[str, list[fx.Node]] = defaultdict(list)
+    # Key by (group_name, transport class); dtype breaks buckets positionally
+    # inside _split_by_dtype_and_size (strict program-adjacency).
+    groups: dict[tuple, list[fx.Node]] = defaultdict(list)
     for node in graph.graph.nodes:
         if not _is_weight_all_gather(node):
             continue
         _, _world, group_name = node.args
-        groups[group_name].append(node)
+        groups[(group_name, split_by(node) if split_by is not None else None)].append(node)
 
     buckets = 0
-    for group_name, ag_nodes in groups.items():
+    for ag_nodes in groups.values():
         for sub in _split_by_dtype_and_size(ag_nodes, node_index, bucket_size_bytes):
             if len(sub) < 2:
                 continue  # single weight -> keep its own all_gather (nothing to coalesce)

magi_compiler/passes/fsdp_overlap/copy_engine.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Graph-side copy-engine transport: bind eligible weights, then retarget tagged gathers.
+
+``bind_weights_for_copy_engine`` runs between lowering and bucketing so a bucket
+stays all-bound or all-unbound. ``rewrite_weight_ag_to_copy_engine`` retargets
+whatever carries ``node_meta.CE_BOUND``. Operators live in ``symm_mem/``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import torch
+import torch.fx as fx
+
+from magi_compiler.utils import magi_logger
+
+from .node_meta import is_ce_bound, is_uneven_shard, is_weight_ag, mark_ce_bound
+
+_ALL_GATHER = torch.ops._c10d_functional.all_gather_into_tensor.default
+_ALL_GATHER_COALESCED = torch.ops._c10d_functional.all_gather_into_tensor_coalesced.default
+
+
+def _weight_holder(src) -> fx.Node | None:
+    """The placeholder/get_attr behind ``to_local(...)``, or None for anything else."""
+    if not isinstance(src, fx.Node) or src.op != "call_method" or src.target != "to_local":
+        return None
+    owner = src.args[0] if src.args else None
+    return owner if isinstance(owner, fx.Node) and owner.op in ("placeholder", "get_attr") else None
+
+
+def copy_engine_weight_candidates(graph: fx.GraphModule) -> list[tuple[fx.Node, fx.Node]]:
+    """Weight all-gathers the copy engine could serve, as ``(gather, holder)`` pairs.
+
+    Holder must be exactly ``to_local(placeholder|get_attr)``: anything in between
+    is a temporary with no peer views. Binding attempts only these.
+    """
+    candidates: list[tuple[fx.Node, fx.Node]] = []
+    for node in graph.graph.nodes:
+        if node.op != "call_function" or node.target is not _ALL_GATHER:
+            continue
+        if not is_weight_ag(node) or is_uneven_shard(node):
+            continue
+        holder = _weight_holder(node.args[0] if node.args else None)
+        if holder is not None:
+            candidates.append((node, holder))
+    return candidates
+
+
+def bind_weights_for_copy_engine(graph: fx.GraphModule, example_inputs: Sequence[Any] | None, min_shard_bytes: int = 0) -> int:
+    """Move eligible weights into symmetric memory and tag the gathers that got served.
+
+    Must run before bucketing. ``example_inputs`` are the live tensors Dynamo
+    captured; without them nothing binds. Returns how many gathers are now
+    copy-engine backed.
+    """
+    from magi_compiler.symm_mem.bind import bind_graph_weights
+
+    placeholders = graph.graph.find_nodes(op="placeholder")
+    served = bind_graph_weights(
+        graph=graph,
+        candidates=copy_engine_weight_candidates(graph),
+        # Built here so the NCCL path never pays for it.
+        placeholder_examples=dict(zip((n.name for n in placeholders), example_inputs or ())),
+        min_shard_bytes=min_shard_bytes,
+    )
+    for node in served:
+        mark_ce_bound(node)
+    return len(served)
+
+
+def rewrite_weight_ag_to_copy_engine(graph: fx.GraphModule) -> int:
+    """Retarget CE_BOUND weight gathers onto copy-engine ops. Returns how many."""
+    from magi_compiler.symm_mem.all_gather import CE_ALL_GATHER, CE_ALL_GATHER_COALESCED
+
+    rewritten = 0
+    skipped = 0
+    for node in graph.graph.nodes:
+        if node.op != "call_function" or node.target not in (_ALL_GATHER, _ALL_GATHER_COALESCED):
+            continue
+        if not is_weight_ag(node):
+            continue
+        if not is_ce_bound(node):
+            skipped += 1
+            continue
+        node.target = CE_ALL_GATHER if node.target is _ALL_GATHER else CE_ALL_GATHER_COALESCED
+        rewritten += 1
+
+    if rewritten:
+        graph.graph.lint()
+        graph.recompile()
+    magi_logger.info(
+        "FSDP copy-engine rewrite: %d weight all-gather(s) retargeted, %d left on NCCL (weight not bound)", rewritten, skipped
+    )
+    return rewritten

magi_compiler/passes/fsdp_overlap/lower_and_bucket.py

@@ -14,15 +14,26 @@
 
 from __future__ import annotations
 
+from typing import Any, Sequence
+
 import torch.fx as fx
 
 from magi_compiler.utils import magi_logger
 
 from .bucket_all_gather import bucket_weight_all_gather_coalesced
+from .copy_engine import bind_weights_for_copy_engine, rewrite_weight_ag_to_copy_engine
+from .node_meta import is_ce_bound
 from .redistribute_lowering import lower_prim_redistribute_to_collectives
 
 
-def lower_and_bucket_full_graph(graph: fx.GraphModule, bucket_mode: str, bucket_size_bytes: int = 0) -> int:
+def lower_and_bucket_full_graph(
+    graph: fx.GraphModule,
+    bucket_mode: str,
+    bucket_size_bytes: int = 0,
+    transport: str = "nccl",
+    example_inputs: Sequence[Any] | None = None,
+    min_shard_bytes: int = 0,
+) -> int:
     """Lower SimpleFSDP weight redistribute -> explicit collectives, then
     optionally bucket them across the WHOLE graph (no subgraph partitioning).
 
@@ -36,19 +47,31 @@ def lower_and_bucket_full_graph(graph: fx.GraphModule, bucket_mode: str, bucket_
     the byte cap in program order (see ``bucket_weight_all_gather_coalesced``).
     0 = no cap (one bucket per (group, dtype) run).
 
+    ``transport="copy_engine"`` wraps the bucketing in the two steps ``copy_engine``
+    owns: binding right after lowering, the retarget at the very end.  Bucketing
+    then keys off what binding served, so bound and unbound gathers are split into
+    separate buckets rather than the unbound ones being dropped from bucketing --
+    losing the copy engine must not also lose coalescing.  ``example_inputs`` and
+    ``min_shard_bytes`` are read only on this path.
+
     Returns the number of buckets created.
     """
     lowered = lower_prim_redistribute_to_collectives(graph)
     magi_logger.info("Whole-graph FSDP lowering: %d weight redistribute -> collectives", lowered)
 
-    bucket_mode = (bucket_mode or "none").lower()
-    if bucket_mode == "none":
-        return 0
+    if transport == "copy_engine":
+        bind_weights_for_copy_engine(graph, example_inputs, min_shard_bytes)
 
+    bucket_mode = (bucket_mode or "none").lower()
+    n = 0
     if bucket_mode == "coalesced":
-        n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes)
-    else:
+        split_by = is_ce_bound if transport == "copy_engine" else None
+        n = bucket_weight_all_gather_coalesced(graph, bucket_size_bytes=bucket_size_bytes, split_by=split_by)
+        magi_logger.info("Whole-graph FSDP bucketing (%s): created %d buckets", bucket_mode, n)
+    elif bucket_mode not in ("none", ""):
         raise ValueError(f"Unknown bucket_mode={bucket_mode!r}; expected 'none' or 'coalesced'")
 
-    magi_logger.info("Whole-graph FSDP bucketing (%s): created %d buckets", bucket_mode, n)
+    if transport == "copy_engine":
+        rewrite_weight_ag_to_copy_engine(graph)
+
     return n