minimal_async_ep: pipeline bounded high-priority copies

tests/unit_tests/test_minimal_async_ep_kernels.py

@@ -19,7 +19,7 @@
 from torch.fx.experimental.proxy_tensor import make_fx
 
 from torchtitan.distributed.minimal_async_ep.kernels import (
-    _copy_rows_to_peer_ptrs_kernel,
+    _copy_rows_to_peer_ptrs_persistent_kernel,
     copy_full_counts_to_peers_kernel,
     copy_rows_to_peers_kernel,
     expand_topk_grad_kernel,
@@ -284,6 +284,33 @@ def test_receive_capacity_rejects_global_routing_overflow():
         )
 
 
+def test_overlap_copy_grid_propagates_to_minimal_async_ep_dispatchers():
+    from torchtitan.experiments.graph_trainer.deepseek_v3.config_registry import (
+        graph_trainer_deepseek_v3_671b_bf16_minimal_async_ep,
+    )
+    from torchtitan.models.common.token_dispatcher import MinimalAsyncEPTokenDispatcher
+
+    config = graph_trainer_deepseek_v3_671b_bf16_minimal_async_ep()
+    config.compile.memory_policy = "full"
+    config.model_spec.model.update_from_config(config=config)
+    dispatchers = [
+        layer.moe.routed_experts.token_dispatcher
+        for layer in config.model_spec.model.layers
+        if layer.moe is not None
+    ]
+    assert dispatchers
+    assert all(dispatcher.num_row_copy_ctas is None for dispatcher in dispatchers)
+
+    config.compile.ep_overlap.enabled = True
+    config.model_spec.model.update_from_config(config=config)
+
+    assert all(
+        isinstance(dispatcher, MinimalAsyncEPTokenDispatcher.Config)
+        and dispatcher.num_row_copy_ctas == 50
+        for dispatcher in dispatchers
+    )
+
+
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
 class TestMinimalAsyncEPKernels(unittest.TestCase):
     @unittest.skipUnless(
@@ -354,6 +381,7 @@ def run_exchange(x, scores, expert_ids):
                 dtype=torch.float32,
                 device=device,
                 receive_capacity_factor=2.0,
+                num_row_copy_ctas=17,
             )
             rank = dist.get_rank()
             x = (
@@ -404,8 +432,32 @@ def wrapped(*args, **kwargs):
 
             buffer_state = minimal_async_ep_api._buffer_state
             assert buffer_state is not None
-            comm_stream = buffer_state.comm_stream.cuda_stream
-            self.assertEqual({stream for _, stream in launch_streams}, {comm_stream})
+            prep_stream = buffer_state.prep_stream.cuda_stream
+            copy_stream = buffer_state.copy_stream.cuda_stream
+            self.assertNotEqual(prep_stream, copy_stream)
+            self.assertLess(
+                buffer_state.copy_stream.priority,
+                buffer_state.prep_stream.priority,
+            )
+            self.assertEqual(
+                buffer_state.num_row_copy_ctas,
+                17,
+            )
+            prep_functions = {
+                "copy_full_counts_to_peers_kernel",
+                "fill_dispatch_metadata_kernel",
+                "fill_combine_metadata_kernel",
+                "invert_flat_indices_kernel",
+            }
+            for name, stream in launch_streams:
+                if name in prep_functions:
+                    self.assertEqual(stream, prep_stream)
+                elif name == "copy_rows_to_peers_kernel":
+                    self.assertEqual(stream, copy_stream)
+            wait_streams = {
+                stream for name, stream in launch_streams if name == "_wait_ready"
+            }
+            self.assertEqual(wait_streams, {prep_stream, copy_stream})
             self.assertEqual(
                 {name for name, _ in launch_streams}, set(launch_functions)
             )
@@ -784,6 +836,26 @@ def test_copy_kernels_handle_strides_and_active_row_masks(self):
         assert_equal(row_dsts[0][4], src[0])
         assert_equal(row_dsts[1][5], torch.zeros(2, device="cuda"))
 
+        bounded_src = torch.arange(128, device="cuda", dtype=torch.float32).view(64, 2)
+        bounded_dst = torch.zeros_like(bounded_src)
+        bounded_dst_ptrs = torch.tensor(
+            [bounded_dst.data_ptr()], device="cuda", dtype=torch.int64
+        )
+        copy_rows_to_peers_kernel(
+            bounded_src,
+            [bounded_dst],
+            torch.zeros(64, device="cuda", dtype=torch.int64),
+            torch.arange(64, device="cuda", dtype=torch.int64),
+            ep_size=1,
+            num_rows=64,
+            num_cols=2,
+            dst_ptrs=bounded_dst_ptrs,
+            block_m=1,
+            num_ctas=2,
+        )
+        torch.cuda.synchronize()
+        assert_equal(bounded_dst, bounded_src)
+
     def test_copy_rows_uses_int64_for_source_stride_arithmetic(self):
         row = 1_048_576
         stride = 2048
@@ -814,7 +886,7 @@ def test_copy_rows_uses_int64_for_source_stride_arithmetic(self):
         src_storage[base_offset + high_offset] = 93
         torch.cuda.synchronize()
 
-        _copy_rows_to_peer_ptrs_kernel[(metadata_numel, 1)](
+        _copy_rows_to_peer_ptrs_persistent_kernel[(20,)](
             src,
             dst_ptrs,
             dst_ranks,
@@ -832,6 +904,7 @@ def test_copy_rows_uses_int64_for_source_stride_arithmetic(self):
             SRC_ROW_DIVISOR=1,
             BLOCK_M=1,
             BLOCK_N=1,
+            NUM_COL_TILES=1,
         )
         torch.cuda.synchronize()
 

torchtitan/distributed/minimal_async_ep/api.py

@@ -92,7 +92,9 @@ class _MinimalAsyncEPBufferState:
     counts_recv_handle: Any
     counts_recv_peer_buffers: list[torch.Tensor]
     counts_recv_peer_ptrs: torch.Tensor
-    comm_stream: torch.cuda.Stream
+    prep_stream: torch.cuda.Stream
+    copy_stream: torch.cuda.Stream
+    num_row_copy_ctas: int | None
     check_receive_capacity: bool
     hidden_recv_buffer_index: int = 0
     pending_events: dict[tuple[str, int], deque[_PendingEvent]] = field(
@@ -184,6 +186,12 @@ def maybe_update_minimal_async_ep_config(model_config: Any, config: Any) -> None
             "MinimalAsyncEP receive capacity factor must be finite and at least "
             "1.0, or None."
         )
+    overlap_enabled = bool(getattr(overlap_config, "enabled", False))
+    num_row_copy_ctas = (
+        getattr(overlap_config, "minimal_async_ep_num_copy_ctas", None)
+        if overlap_enabled
+        else None
+    )
 
     for token_dispatcher_cfg in dispatcher_cfgs:
         token_dispatcher_cfg.hidden_dim = model_config.dim
@@ -193,6 +201,7 @@ def maybe_update_minimal_async_ep_config(model_config: Any, config: Any) -> None
         token_dispatcher_cfg.dtype = TORCH_DTYPE_MAP[
             config.training.mixed_precision_param
         ]
+        token_dispatcher_cfg.num_row_copy_ctas = num_row_copy_ctas
         token_dispatcher_cfg.force_load_balance = config.debug.moe_force_load_balance
         token_dispatcher_cfg.receive_capacity_factor = receive_capacity_factor
         token_dispatcher_cfg.receive_capacity = _buffer_rows(
@@ -285,6 +294,7 @@ def init_buffer(
     device: torch.device,
     force_load_balance: bool = False,
     receive_capacity_factor: float | None = None,
+    num_row_copy_ctas: int | None = None,
 ) -> None:
     """Initialize the process-local MinimalAsyncEP symmetric-memory buffer."""
     global _buffer_state
@@ -294,6 +304,11 @@ def init_buffer(
     num_experts = ep_size * num_local_experts
     assert _buffer_state is None
 
+    if num_row_copy_ctas is not None and num_row_copy_ctas < 1:
+        raise ValueError(
+            f"num_row_copy_ctas must be positive or None, got {num_row_copy_ctas}."
+        )
+
     dispatch_rows, combine_rows = _buffer_rows(
         tokens_per_rank,
         top_k,
@@ -316,14 +331,15 @@ def init_buffer(
     logger.info(
         "Initializing MinimalAsyncEP buffer: hidden_dim=%d, tokens_per_rank=%d, "
         "top_k=%d, num_local_experts=%d, ep_size=%d, capacity_mode=%s, "
-        "buffer_rows=%d",
+        "buffer_rows=%d, num_row_copy_ctas=%s",
         hidden_dim,
         tokens_per_rank,
         top_k,
         num_local_experts,
         ep_size,
         capacity_mode,
         dispatch_rows,
+        num_row_copy_ctas if num_row_copy_ctas is not None else "unbounded",
     )
     backend = symm_mem.get_backend(device)
     if backend != "CUDA":
@@ -386,7 +402,9 @@ def init_buffer(
         counts_recv_handle=counts_recv_handle,
         counts_recv_peer_buffers=counts_recv_peer_buffers,
         counts_recv_peer_ptrs=counts_recv_peer_ptrs,
-        comm_stream=torch.cuda.Stream(device=device),
+        prep_stream=torch.cuda.Stream(device=device),
+        copy_stream=torch.cuda.Stream(device=device, priority=-1),
+        num_row_copy_ctas=num_row_copy_ctas,
         check_receive_capacity=(
             receive_capacity_factor is not None and not force_load_balance
         ),
@@ -442,7 +460,7 @@ def _copy_rows_to_peers_async_cuda(
     num_valid_rows: torch.Tensor | None = None,
     retained: tuple[torch.Tensor, ...] = (),
 ) -> torch.Tensor:
-    """Launch a row copy through symmetric memory on the EP comm stream."""
+    """Launch a row copy through symmetric memory on the EP copy stream."""
     assert _buffer_state is not None
 
     buffer_rows = _buffer_state.hidden_recv_buffers[0].shape[0]
@@ -468,9 +486,9 @@ def _copy_rows_to_peers_async_cuda(
     )
 
     launch_stream = torch.cuda.current_stream(x.device)
-    comm_stream = _buffer_state.comm_stream
-    comm_stream.wait_stream(launch_stream)
-    with torch.cuda.stream(comm_stream):
+    copy_stream = _buffer_state.copy_stream
+    copy_stream.wait_stream(launch_stream)
+    with torch.cuda.stream(copy_stream):
         copy_rows_to_peers_kernel(
             x,
             hidden_recv_peer_buffers,
@@ -485,9 +503,10 @@ def _copy_rows_to_peers_async_cuda(
             src_row_divisor=src_row_divisor,
             dst_ptrs=hidden_recv_peer_ptrs,
             num_valid_rows=num_valid_rows,
+            num_ctas=_buffer_state.num_row_copy_ctas,
         )
         _wait_ready(hidden_recv_handle, _HIDDEN_READY_CHANNEL)
-        _record_pending_event(exchange, hidden_recv_view, comm_stream, retained)
+        _record_pending_event(exchange, hidden_recv_view, copy_stream, retained)
     return hidden_recv_view
 
 
@@ -764,9 +783,9 @@ def dispatch_op(
     )
 
     launch_stream = torch.cuda.current_stream(dispatch_input.device)
-    comm_stream = _buffer_state.comm_stream
-    comm_stream.wait_stream(launch_stream)
-    with torch.cuda.stream(comm_stream):
+    prep_stream = _buffer_state.prep_stream
+    prep_stream.wait_stream(launch_stream)
+    with torch.cuda.stream(prep_stream):
         torch.argsort(
             T_row_to_expert_N,
             stable=True,

torchtitan/distributed/minimal_async_ep/kernels.py

@@ -351,6 +351,64 @@ def _copy_rows_to_peer_ptrs_kernel(
     tl.store(dst_ptr, values, mask=mask & dst_rank_mask[:, None])
 
 
+# Preserve the 2-D launch above for unbounded copies. Flattening its grid changes
+# peer issue order and causes severe rank imbalance at EP64.
+@triton.jit
+def _copy_rows_to_peer_ptrs_persistent_kernel(
+    src,
+    dst_ptrs: tl.pointer_type(tl.int64),
+    dst_ranks: tl.pointer_type(tl.int64),
+    dst_rows: tl.pointer_type(tl.int64),
+    num_valid_rows: tl.pointer_type(tl.int64),
+    src_rows: tl.pointer_type(tl.int64),
+    NUM_ROWS: tl.constexpr,
+    NUM_COLS: tl.constexpr,
+    SRC_ROW_STRIDE: tl.constexpr,
+    SRC_COL_STRIDE: tl.constexpr,
+    DST_ROW_STRIDE: tl.constexpr,
+    DST_DTYPE: tl.constexpr,
+    HAS_NUM_VALID_ROWS: tl.constexpr,
+    HAS_SRC_ROWS: tl.constexpr,
+    SRC_ROW_DIVISOR: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    NUM_COL_TILES: tl.constexpr,
+) -> None:
+    row_limit = NUM_ROWS
+    if HAS_NUM_VALID_ROWS:
+        row_limit = tl.load(num_valid_rows)
+    num_tiles = tl.cdiv(row_limit, BLOCK_M) * NUM_COL_TILES
+    tile = tl.program_id(0)
+    while tile < num_tiles:
+        row_tile = tile // NUM_COL_TILES
+        col_tile = tile % NUM_COL_TILES
+        row = (row_tile * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64)
+        col = col_tile * BLOCK_N + tl.arange(0, BLOCK_N)
+        row_mask = row < row_limit
+        col_mask = col < NUM_COLS
+        mask = row_mask[:, None] & col_mask[None, :]
+        src_row = row
+        if HAS_SRC_ROWS:
+            src_row = tl.load(src_rows + row, mask=row_mask, other=0).to(tl.int64)
+            if SRC_ROW_DIVISOR != 1:
+                src_row = src_row // SRC_ROW_DIVISOR
+        dst_rank = tl.load(dst_ranks + row, mask=row_mask, other=-1)
+        dst_row = tl.load(dst_rows + row, mask=row_mask, other=0).to(tl.int64)
+        dst_rank_mask = row_mask & (dst_rank >= 0)
+        values = tl.load(
+            src + src_row[:, None] * SRC_ROW_STRIDE + col[None, :] * SRC_COL_STRIDE,
+            mask=mask,
+        )
+        dst_base = tl.load(dst_ptrs + dst_rank, mask=dst_rank_mask, other=0)
+        dst_ptr = (
+            dst_base.to(tl.pointer_type(DST_DTYPE))[:, None]
+            + dst_row[:, None] * DST_ROW_STRIDE
+            + col[None, :]
+        )
+        tl.store(dst_ptr, values, mask=mask & dst_rank_mask[:, None])
+        tile += tl.num_programs(0)
+
+
 def copy_full_counts_to_peers_kernel(
     counts: torch.Tensor,
     dsts: list[torch.Tensor],
@@ -397,22 +455,29 @@ def copy_rows_to_peers_kernel(
     src_rows: torch.Tensor | None = None,
     src_row_divisor: int = 1,
     num_valid_rows: torch.Tensor | None = None,
+    num_ctas: int | None = None,
 ) -> None:
     if len(dsts) != ep_size:
         raise ValueError(f"expected {ep_size} destination buffers, got {len(dsts)}.")
 
+    if num_ctas is not None and num_ctas <= 0:
+        raise ValueError(f"num_ctas must be positive, got {num_ctas}.")
+
     block_n = min(_MAX_BLOCK_N, triton.next_power_of_2(num_cols))
-    grid = (triton.cdiv(num_rows, block_m), triton.cdiv(num_cols, block_n))
+    num_row_tiles = triton.cdiv(num_rows, block_m)
+    num_col_tiles = triton.cdiv(num_cols, block_n)
     dst_dtype = _HIDDEN_ROW_DTYPES.get(src.dtype)
     if dst_dtype is None:
         raise ValueError(f"Unsupported MinimalAsyncEP row-copy dtype: {src.dtype}.")
-    _copy_rows_to_peer_ptrs_kernel[grid](
+    common_args = (
         src,
         dst_ptrs,
         dst_ranks,
         dst_rows,
         num_valid_rows if num_valid_rows is not None else dst_rows[:1],
         src_rows if src_rows is not None else dst_rows,
+    )
+    common_kwargs = dict(
         NUM_ROWS=num_rows,
         NUM_COLS=num_cols,
         SRC_ROW_STRIDE=src.stride(0),
@@ -426,6 +491,18 @@ def copy_rows_to_peers_kernel(
         BLOCK_N=block_n,
         num_warps=num_warps,
     )
+    if num_ctas is None:
+        _copy_rows_to_peer_ptrs_kernel[(num_row_tiles, num_col_tiles)](
+            *common_args,
+            **common_kwargs,
+        )
+    else:
+        num_tiles = num_row_tiles * num_col_tiles
+        _copy_rows_to_peer_ptrs_persistent_kernel[(min(num_tiles, num_ctas),)](
+            *common_args,
+            NUM_COL_TILES=num_col_tiles,
+            **common_kwargs,
+        )
 
 
 def fill_dispatch_metadata_kernel(