From b841bb0fd740cab990240f31510989fac233a1e4 Mon Sep 17 00:00:00 2001 From: Sanket Jayant Purandare Date: Wed, 29 Jul 2026 16:40:05 -0700 Subject: [PATCH] minimal_async_ep: pipeline bounded high-priority copies ## Summary Run dispatch preparation on a normal stream, then hand row copies to a high-priority stream. Bound overlapping copies with a persistent CTA grid so expert compute retains execution resources; keep the no-overlap single-buffer path unbounded because it lies on the critical path. Expose the grid as `--compile.ep_overlap.minimal_async_ep_num_copy_ctas`. The default is the measured 50-CTA GB300 setting, while `None` selects an unbounded overlap control for hardware-specific tuning. ## Why Moving copies to another stream does not reserve SM capacity for expert work. A bounded persistent launch provides an explicit resource limit while stream priority minimizes gaps between dispatch preparation and row copies. ## Test Plan - Verify CTA-grid propagation into all MinimalAsyncEP dispatchers. - Reject invalid configured grid sizes. - Exercise bounded copy correctness and prep/copy stream placement in the two-rank dispatcher test. [ghstack-poisoned] --- .../test_minimal_async_ep_kernels.py | 79 ++++++++++++++++++- .../distributed/minimal_async_ep/api.py | 43 +++++++--- .../distributed/minimal_async_ep/kernels.py | 68 +++++++++------- .../experiments/graph_trainer/configs.py | 14 ++++ .../graph_trainer/tests/test_passes.py | 7 ++ torchtitan/models/common/token_dispatcher.py | 6 ++ 6 files changed, 175 insertions(+), 42 deletions(-) diff --git a/tests/unit_tests/test_minimal_async_ep_kernels.py b/tests/unit_tests/test_minimal_async_ep_kernels.py index e9fe72f35f..0eaef660b0 100644 --- a/tests/unit_tests/test_minimal_async_ep_kernels.py +++ b/tests/unit_tests/test_minimal_async_ep_kernels.py @@ -237,6 +237,33 @@ def test_671b_forced_load_balance_uses_tight_receive_capacity(): assert dispatcher.receive_capacity == 131_072 +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( @@ -306,6 +333,7 @@ def run_exchange(x, scores, expert_ids): top_k=2, dtype=torch.float32, device=device, + num_row_copy_ctas=17, ) rank = dist.get_rank() x = ( @@ -356,8 +384,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) ) @@ -736,6 +788,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 @@ -766,7 +838,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_kernel[(20,)]( src, dst_ptrs, dst_ranks, @@ -784,6 +856,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() diff --git a/torchtitan/distributed/minimal_async_ep/api.py b/torchtitan/distributed/minimal_async_ep/api.py index 0179c37dd1..54fcd36cb2 100644 --- a/torchtitan/distributed/minimal_async_ep/api.py +++ b/torchtitan/distributed/minimal_async_ep/api.py @@ -91,7 +91,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 hidden_recv_buffer_index: int = 0 pending_events: dict[tuple[str, int], deque[_PendingEvent]] = field( default_factory=dict @@ -169,6 +171,14 @@ def maybe_update_minimal_async_ep_config(model_config: Any, config: Any) -> None "--compile.memory_policy full for graph_trainer." ) + overlap_config = getattr(config.compile, "ep_overlap", 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 token_dispatcher_cfg.tokens_per_rank = ( @@ -177,6 +187,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 = _buffer_rows( token_dispatcher_cfg.tokens_per_rank, @@ -261,6 +272,7 @@ def init_buffer( dtype: torch.dtype, device: torch.device, force_load_balance: bool = False, + num_row_copy_ctas: int | None = None, ) -> None: """Initialize the process-local MinimalAsyncEP symmetric-memory buffer.""" global _buffer_state @@ -270,6 +282,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, @@ -283,7 +300,7 @@ 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, @@ -291,6 +308,7 @@ def init_buffer( 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": @@ -353,7 +371,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, ) @@ -406,7 +426,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] @@ -432,9 +452,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, @@ -449,9 +469,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 @@ -716,9 +737,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, diff --git a/torchtitan/distributed/minimal_async_ep/kernels.py b/torchtitan/distributed/minimal_async_ep/kernels.py index 086b31b3df..25db0370bd 100644 --- a/torchtitan/distributed/minimal_async_ep/kernels.py +++ b/torchtitan/distributed/minimal_async_ep/kernels.py @@ -307,6 +307,7 @@ def _copy_rows_to_peer_ptrs_kernel( SRC_ROW_DIVISOR: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + NUM_COL_TILES: tl.constexpr, ) -> None: """Copy rows into peer symmetric hidden buffers through pointer tables. @@ -319,36 +320,39 @@ def _copy_rows_to_peer_ptrs_kernel( ``dst_rows=[3, 4]``, and ``src_rows=[2, 0]``, peer 1 row 3 receives ``[30]`` and peer 0 row 4 receives ``[10]``. """ - row_start = tl.program_id(0) * BLOCK_M row_limit = NUM_ROWS if HAS_NUM_VALID_ROWS: row_limit = tl.load(num_valid_rows) - if row_start >= row_limit: - return - row = (row_start + tl.arange(0, BLOCK_M)).to(tl.int64) - col = tl.program_id(1) * 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]) + 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( @@ -397,12 +401,19 @@ 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) + num_tiles = num_row_tiles * num_col_tiles + grid = (min(num_tiles, num_ctas) if num_ctas is not None else num_tiles,) dst_dtype = _HIDDEN_ROW_DTYPES.get(src.dtype) if dst_dtype is None: raise ValueError(f"Unsupported MinimalAsyncEP row-copy dtype: {src.dtype}.") @@ -424,6 +435,7 @@ def copy_rows_to_peers_kernel( SRC_ROW_DIVISOR=src_row_divisor, BLOCK_M=block_m, BLOCK_N=block_n, + NUM_COL_TILES=num_col_tiles, num_warps=num_warps, ) diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 4a3fcbbf34..4ac6a4b8b4 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -52,6 +52,13 @@ class EpOverlapConfig: produced by either eager or graph chunking. """ + minimal_async_ep_num_copy_ctas: int | None = 50 + """Persistent row-copy grid size used by MinimalAsyncEP during overlap. + + ``None`` leaves the copy kernel unbounded. This setting has no effect on + other EP backends or when EP overlap is disabled. + """ + disable_early_grad_accumulation: bool = False """Disable chunked parameter-gradient accumulation before communication. @@ -207,6 +214,13 @@ def validate_ep_overlap_config( "--compile.ep_overlap.module_fqn layers.*.moe" ) + num_copy_ctas = ep_overlap_config.minimal_async_ep_num_copy_ctas + if num_copy_ctas is not None and num_copy_ctas < 1: + raise ValueError( + "--compile.ep_overlap.minimal_async_ep_num_copy_ctas must be " + "positive or None" + ) + return chunk_dim, chunk_strategy, module_fqn diff --git a/torchtitan/experiments/graph_trainer/tests/test_passes.py b/torchtitan/experiments/graph_trainer/tests/test_passes.py index 3103666e7a..cd2cc5387e 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_passes.py @@ -39,6 +39,7 @@ from torchtitan.experiments.graph_trainer.configs import ( EpOverlapConfig, GraphTrainerCompileConfig, + validate_ep_overlap_config, ) from torchtitan.experiments.graph_trainer.cudagraph import ( insert_kernel_annotations_pass, @@ -3618,6 +3619,12 @@ def test_ep_overlap_pass_pipeline_order(self): disabled_names.index("joint_transformer_block_bucketing_reordering_pass"), ) + def test_ep_overlap_rejects_invalid_minimal_async_ep_copy_grid(self): + with self.assertRaisesRegex(ValueError, "must be positive or None"): + validate_ep_overlap_config( + EpOverlapConfig(enabled=True, minimal_async_ep_num_copy_ctas=0) + ) + def test_graph_ep_chunking_rejects_tensor_parallel(self): cases = ( ("seq", "layers.*.moe"), diff --git a/torchtitan/models/common/token_dispatcher.py b/torchtitan/models/common/token_dispatcher.py index 83e4560851..7c94c16f95 100644 --- a/torchtitan/models/common/token_dispatcher.py +++ b/torchtitan/models/common/token_dispatcher.py @@ -1114,6 +1114,7 @@ class MinimalAsyncEPTokenDispatcher(LocalTokenDispatcher): tokens_per_rank: int | None dtype: torch.dtype | None buffer_device: torch.device + num_row_copy_ctas: int | None force_load_balance: bool receive_capacity: int | None @@ -1123,6 +1124,7 @@ class Config(LocalTokenDispatcher.Config): tokens_per_rank: int | None = None dtype: torch.dtype | None = None device: torch.device | None = None + num_row_copy_ctas: int | None = None force_load_balance: bool = False receive_capacity: int | None = None @@ -1133,6 +1135,7 @@ def __init__(self, config: Config): self.hidden_dim = config.hidden_dim self.tokens_per_rank = config.tokens_per_rank self.dtype = config.dtype + self.num_row_copy_ctas = config.num_row_copy_ctas self.force_load_balance = config.force_load_balance self.receive_capacity = config.receive_capacity if config.device is None: @@ -1154,6 +1157,7 @@ def __init__(self, config: Config): int, torch.dtype, torch.device, + int | None, bool, int, ] @@ -1217,6 +1221,7 @@ def init_buffer(self) -> None: self.top_k, self.dtype, self.buffer_device, + self.num_row_copy_ctas, self.force_load_balance, self.receive_capacity, ) @@ -1237,6 +1242,7 @@ def init_buffer(self) -> None: dtype=self.dtype, device=self.buffer_device, force_load_balance=self.force_load_balance, + num_row_copy_ctas=self.num_row_copy_ctas, ) MinimalAsyncEPTokenDispatcher._global_buffer_key = buffer_key