diff --git a/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md index 57f7f4fe..d6276e9f 100644 --- a/python/mscclpp/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -69,8 +69,7 @@ class MoECommunicatorConfig: # Quantization defaults quant: Optional[QuantConfig] = None - # Transport resources - num_rdma_qps_per_rank: int = 12 # RDMA QPs per peer rank; advanced tuning + # Launch resources num_sms: int = 20 # Overlap @@ -141,7 +140,7 @@ a later version can add an explicit `expert_map` for arbitrary placement. | `max_tokens_per_rank` | dispatch capacity | | `max_recv_tokens_per_rank` | recv buffer capacity | | scratch buffers | internally sized from mode, capacity, topology, and shape | -| `num_rdma_qps_per_rank`, `num_sms` | backend launch/resource tuning | +| `num_sms` | backend launch/resource tuning | | `dispatch_config`, `combine_config` | backend-specific tuning configs | | `overlap_capability` | whether selected MLP/backend supports notify | @@ -152,8 +151,9 @@ specialized advanced path. The active implementation supports `mode=MoEMode.LOW_LATENCY` and `mode=MoEMode.HIGH_THROUGHPUT`. `mode` must be a `MoEMode` enum value, not a -string. LL uses an expert-major output layout; HT uses a flat output layout and -selects intranode vs internode transport from the runtime size hints. +string. LL uses an expert-major output layout. HT uses a flat output layout and +supports only 2, 4, or 8 ranks within one detected GPU IPC/NVL fabric domain; +that domain may span multiple hosts. ```python moe_comm = MoECommunicator(..., mode=MoEMode.LOW_LATENCY) @@ -304,16 +304,11 @@ class ExpertMajorCombineContext: @dataclass -class RowMajorIntranodeCombineContext: +class RowMajorCombineContext: ... -@dataclass -class RowMajorInternodeCombineContext: - ... - - -CombineContext = ExpertMajorCombineContext | RowMajorIntranodeCombineContext | RowMajorInternodeCombineContext +CombineContext = ExpertMajorCombineContext | RowMajorCombineContext class DispatchHandle: @@ -326,12 +321,8 @@ class ExpertMajorDispatchHandle(DispatchHandle): combine_context: ExpertMajorCombineContext -class RowMajorIntranodeDispatchHandle(DispatchHandle): - combine_context: RowMajorIntranodeCombineContext - - -class RowMajorInternodeDispatchHandle(DispatchHandle): - combine_context: RowMajorInternodeCombineContext +class RowMajorDispatchHandle(DispatchHandle): + combine_context: RowMajorCombineContext @dataclass @@ -418,7 +409,7 @@ overlap is operation-level only. Each concrete `DispatchHandle` stores a layout-specific `combine_context` used to reverse dispatch and finish combine. `ExpertMajorDispatchHandle` uses `ExpertMajorCombineContext` (`topk_ids`, `weights`, source info, layout ranges, -shape, and capacity). Row-major handles use intranode or internode combine contexts with +shape, and capacity). Row-major handles use the intranode combine context with receive-side weights, source indices, prefix matrices, and send-head tensors. The MLP should treat the handle as opaque and pass it back to `combine`. diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py index 1aa7f15d..2719de2f 100644 --- a/python/mscclpp/ep/__init__.py +++ b/python/mscclpp/ep/__init__.py @@ -27,10 +27,8 @@ from .communicator import ( # noqa: F401 MoEMode, OperationOverlapConfig, QuantConfig, - RowMajorInternodeDispatchHandle, - RowMajorInternodeCombineContext, - RowMajorIntranodeDispatchHandle, - RowMajorIntranodeCombineContext, + RowMajorDispatchHandle, + RowMajorCombineContext, ) __all__ = [ @@ -51,8 +49,6 @@ __all__ = [ "MoEMode", "OperationOverlapConfig", "QuantConfig", - "RowMajorInternodeDispatchHandle", - "RowMajorInternodeCombineContext", - "RowMajorIntranodeDispatchHandle", - "RowMajorIntranodeCombineContext", + "RowMajorDispatchHandle", + "RowMajorCombineContext", ] diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py index ffa18de6..fca7f2c4 100644 --- a/python/mscclpp/ep/communicator.py +++ b/python/mscclpp/ep/communicator.py @@ -24,10 +24,8 @@ from .types import ( MoECommunicatorConfig, OperationOverlapConfig, QuantConfig, - RowMajorInternodeDispatchHandle, - RowMajorInternodeCombineContext, - RowMajorIntranodeDispatchHandle, - RowMajorIntranodeCombineContext, + RowMajorDispatchHandle, + RowMajorCombineContext, ) __all__ = [ @@ -48,10 +46,8 @@ __all__ = [ "MoEMode", "OperationOverlapConfig", "QuantConfig", - "RowMajorInternodeDispatchHandle", - "RowMajorInternodeCombineContext", - "RowMajorIntranodeDispatchHandle", - "RowMajorIntranodeCombineContext", + "RowMajorDispatchHandle", + "RowMajorCombineContext", ] diff --git a/python/mscclpp/ep/high_throughput.py b/python/mscclpp/ep/high_throughput.py index 081d61e5..9f025d86 100644 --- a/python/mscclpp/ep/high_throughput.py +++ b/python/mscclpp/ep/high_throughput.py @@ -3,32 +3,13 @@ # # Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), # branch ``chhwang/dev-atomic-add-cleanup``. Licensed under the MIT License. -"""High-throughput backend for the high-level MoE communicator. +"""Fabric-domain high-throughput backend for the high-level MoE communicator. -This module contains both the high-level HT backend used by -``MoECommunicator`` and the raw-pointer wrapper around the nanobind extension -``mscclpp_ep_cpp.ExpertParallelRuntime`` (the DeepEP-style high-throughput -runtime). The extension exposes a **torch-free, raw-pointer** boundary identical -in spirit to the low-latency ``MoERuntime``: every device tensor crosses the -boundary as an integer ``tensor.data_ptr()`` plus explicit shape/size arguments, -so the module never links libtorch. - -Because the C++ side no longer allocates the data-dependent receive buffers, -dynamic recv sizing uses an explicit **two-phase** protocol on the intranode / -internode dispatch path: - -1. ``*_notify_dispatch`` runs the size-exchange kernel and returns - ``num_recv_tokens`` (and, internode, ``num_rdma_recv_tokens``), writing the - routing prefix matrices into caller-provided tensors. -2. The wrapper allocates the recv output tensors sized by ``num_recv_tokens`` - (or, for the zero-copy direct path, views this rank's recv pool via - :meth:`resolve_intranode_recv_x_buffer`). -3. ``*_dispatch`` runs the data-movement kernel into those output pointers. - -The cached fast path skips the notify phase (``cached_mode=True``) by reusing a -previous dispatch's prefix matrices and recv count. - -The low-latency path is served by ``low_latency.py``. +The C++ runtime follows the low-latency resource model: it reuses the existing +MSCCL++ communicator, owns one physical symmetric buffer, and exposes a +torch-free raw-pointer boundary. Dynamic receive sizing uses a two-phase +``notify_dispatch`` then ``dispatch`` protocol. Cached dispatches reuse the +previous routing matrices and receive count. """ from __future__ import annotations @@ -43,17 +24,13 @@ from .types import ( DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, - RowMajorInternodeDispatchHandle, - RowMajorInternodeCombineContext, - RowMajorIntranodeDispatchHandle, - RowMajorIntranodeCombineContext, MoECommunicatorConfig, QuantConfig, + RowMajorCombineContext, + RowMajorDispatchHandle, ) from .utils import ( - all_gather_object as _all_gather_object, bf16_view as _bf16_view, - broadcast_object as _broadcast_object, current_stream_ptr as _stream_ptr, exclusive_cumsum, ptr as _ptr, @@ -64,9 +41,8 @@ from .utils import ( class HighThroughputRuntime: """Core high-throughput expert-parallel (EP) communication runtime. - ``comm`` is the ``mscclpp.CommGroup`` used for rank information and - out-of-band exchange of device ids, CUDA-IPC handles, and the MSCCL++ unique - id. All dispatch/combine data movement happens through the MSCCL++ runtime. + ``comm`` provides the initialized MSCCL++ communicator used to exchange and + map the intranode physical symmetric buffers. """ #: Default number of SMs reserved for comms kernels. Matches DeepEP. @@ -75,43 +51,13 @@ class HighThroughputRuntime: def __init__( self, comm: Any, - num_nvl_bytes: int = 0, - num_rdma_bytes: int = 0, - low_latency_mode: bool = False, - num_qps_per_rank: int = 12, + max_hidden_bytes: int, + config: Config, ) -> None: - if low_latency_mode: - raise NotImplementedError( - "HighThroughputRuntime serves the high-throughput path only; use MoERuntime for low latency." - ) - if num_qps_per_rank <= 0: - raise ValueError("num_qps_per_rank must be > 0") - self.rank: int = comm.my_rank self.group_size: int = comm.nranks self.comm = comm - self.num_nvl_bytes = num_nvl_bytes - self.num_rdma_bytes = num_rdma_bytes - self.num_qps_per_rank = num_qps_per_rank - - self.runtime = _cpp.ExpertParallelRuntime(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes) - - # Exchange device ids + CUDA-IPC handles + (for RDMA) the MSCCL++ unique id. - local_device_id = self.runtime.get_local_device_id() - device_ids = _all_gather_object(comm, local_device_id, 0xE000) - - local_ipc_handle = self.runtime.get_local_ipc_handle() - ipc_handles = _all_gather_object(comm, local_ipc_handle, 0xE100) - - root_unique_id: Optional[bytes] = None - if self.rank == 0: - root_unique_id = self.runtime.create_unique_id() - root_unique_id = _broadcast_object(comm, root_unique_id, 0, 0xE200) - assert root_unique_id is not None - self.runtime.connect(root_unique_id) - - ipc_handles_ba = [bytearray(h) if h is not None else None for h in ipc_handles] - self.runtime.sync(device_ids, ipc_handles_ba, bytearray(root_unique_id)) + self.runtime = _cpp.ExpertParallelRuntime(comm.communicator, max_hidden_bytes, config) # ------------------------------------------------------------------ # Sanity helpers @@ -123,40 +69,21 @@ class HighThroughputRuntime: def is_internode_available(self) -> bool: return self.runtime.is_internode_available() - def get_local_device_id(self) -> int: - return self.runtime.get_local_device_id() - - def get_num_rdma_ranks(self) -> int: - return self.runtime.get_num_rdma_ranks() - - def get_rdma_rank(self) -> int: - return self.runtime.get_rdma_rank() - - def get_root_rdma_rank(self, global_: bool) -> int: - return self.runtime.get_root_rdma_rank(global_) - # ------------------------------------------------------------------ # Dispatch layout # ------------------------------------------------------------------ def get_dispatch_layout(self, topk_idx: torch.Tensor, num_experts: int): - """Returns ``(num_tokens_per_rank, num_tokens_per_rdma_rank|None, - num_tokens_per_expert, is_token_in_rank)``.""" + """Return per-rank, per-expert, and token-membership routing metadata.""" assert topk_idx.dim() == 2 and topk_idx.is_contiguous() num_tokens, num_topk = int(topk_idx.size(0)), int(topk_idx.size(1)) num_tokens_per_rank = torch.empty((self.group_size,), dtype=torch.int32, device="cuda") num_tokens_per_expert = torch.empty((num_experts,), dtype=torch.int32, device="cuda") is_token_in_rank = torch.empty((num_tokens, self.group_size), dtype=torch.bool, device="cuda") - num_tokens_per_rdma_rank = None - if self.is_internode_available(): - num_tokens_per_rdma_rank = torch.empty( - (self.runtime.get_num_rdma_ranks(),), dtype=torch.int32, device="cuda" - ) - self.runtime.get_dispatch_layout( + self.runtime.layout( _ptr(num_tokens_per_rank), - _ptr(num_tokens_per_rdma_rank), _ptr(num_tokens_per_expert), _ptr(is_token_in_rank), _ptr(topk_idx), @@ -165,13 +92,13 @@ class HighThroughputRuntime: num_experts, _stream_ptr(), ) - return num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank + return num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank # ------------------------------------------------------------------ - # Intranode dispatch (two-phase) + combine + # Dispatch (two-phase) + combine # ------------------------------------------------------------------ - def intranode_dispatch( + def dispatch( self, x: torch.Tensor, x_scales: Optional[torch.Tensor], @@ -184,19 +111,13 @@ class HighThroughputRuntime: cached_rank_prefix_matrix: Optional[torch.Tensor], cached_channel_prefix_matrix: Optional[torch.Tensor], expert_alignment: int, - config, ): - """High-throughput intranode dispatch. - - Returns ``(recv_x, recv_x_scales, recv_topk_idx, recv_topk_weights, - num_recv_tokens_per_expert_list, rank_prefix_matrix, - channel_prefix_matrix, recv_channel_prefix_matrix, recv_src_idx, - send_head)`` to mirror the previous DeepEP-style surface.""" + """Run high-throughput dispatch and return outputs plus combine metadata.""" assert x.dim() == 2 and x.is_contiguous() cached_mode = cached_rank_prefix_matrix is not None num_tokens, hidden = int(x.size(0)), int(x.size(1)) x_element_size = x.element_size() - num_channels = self.runtime.get_intranode_dispatch_num_channels(x_element_size, config) + num_channels = self.runtime.get_dispatch_num_channels(x_element_size) num_topk = int(topk_idx.size(1)) if topk_idx is not None else 0 num_scales = 0 @@ -217,7 +138,7 @@ class HighThroughputRuntime: rank_prefix_matrix = torch.empty((self.group_size, self.group_size), dtype=torch.int32, device="cuda") channel_prefix_matrix = torch.empty((self.group_size, num_channels), dtype=torch.int32, device="cuda") num_recv_per_expert_host = torch.empty((num_local_experts,), dtype=torch.int32, device="cpu") - num_recv_tokens = self.runtime.intranode_notify_dispatch( + num_recv_tokens = self.runtime.notify_dispatch( _ptr(rank_prefix_matrix), _ptr(channel_prefix_matrix), _ptr(num_recv_per_expert_host), @@ -228,13 +149,12 @@ class HighThroughputRuntime: num_experts, x_element_size, expert_alignment, - config, _stream_ptr(), ) num_recv_tokens_per_expert_list = num_recv_per_expert_host.tolist() # ----- Phase B: allocate recv outputs (or view the recv pool) ----- - recv_x = self._alloc_recv_x(num_recv_tokens, hidden, x_element_size, config) + recv_x = self._alloc_recv_x(num_tokens, num_recv_tokens, hidden, x_element_size) recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int32, device="cuda") send_head = torch.empty((num_tokens, self.group_size), dtype=torch.int32, device="cuda") recv_channel_prefix_matrix = torch.empty((self.group_size, num_channels), dtype=torch.int32, device="cuda") @@ -252,7 +172,7 @@ class HighThroughputRuntime: else None ) - self.runtime.intranode_dispatch( + self.runtime.dispatch( _ptr(recv_x), _ptr(recv_x_scales), _ptr(recv_topk_idx), @@ -275,7 +195,6 @@ class HighThroughputRuntime: x_element_size, num_recv_tokens, cached_mode, - config, _stream_ptr(), ) return ( @@ -291,16 +210,16 @@ class HighThroughputRuntime: send_head, ) - def _alloc_recv_x(self, num_recv_tokens: int, hidden: int, x_element_size: int, config) -> torch.Tensor: + def _alloc_recv_x(self, num_tokens: int, num_recv_tokens: int, hidden: int, x_element_size: int) -> torch.Tensor: """Allocate ``recv_x`` or, when the zero-copy direct path is active, view this rank's recv pool (so the sender writes hidden straight to its final slot and the TMA combine gathers from the same pool).""" - pool_ptr = self.runtime.resolve_intranode_recv_x_buffer(num_recv_tokens, hidden, x_element_size, config) + pool_ptr = self.runtime.resolve_recv_x_buffer(num_tokens, num_recv_tokens, hidden, x_element_size) if pool_ptr != 0: return _bf16_view(pool_ptr, num_recv_tokens, hidden, owner=self) return torch.empty((num_recv_tokens, hidden), dtype=torch.bfloat16, device="cuda") - def intranode_combine( + def combine( self, x: torch.Tensor, topk_weights: Optional[torch.Tensor], @@ -308,7 +227,6 @@ class HighThroughputRuntime: rank_prefix_matrix: torch.Tensor, channel_prefix_matrix: torch.Tensor, send_head: torch.Tensor, - config, ): """Returns ``(combined_x, combined_topk_weights|None)``.""" assert x.dim() == 2 and x.is_contiguous() @@ -323,7 +241,7 @@ class HighThroughputRuntime: if topk_weights is not None else None ) - self.runtime.intranode_combine( + self.runtime.combine( _ptr(combined_x), _ptr(combined_topk_weights), _ptr(x), @@ -338,243 +256,6 @@ class HighThroughputRuntime: num_topk, x.element_size(), ring_num_channels, - config, - _stream_ptr(), - ) - return combined_x, combined_topk_weights - - # ------------------------------------------------------------------ - # Internode dispatch (two-phase) + combine - # ------------------------------------------------------------------ - - def internode_dispatch( - self, - x: torch.Tensor, - x_scales: Optional[torch.Tensor], - topk_idx: Optional[torch.Tensor], - topk_weights: Optional[torch.Tensor], - num_tokens_per_rank: Optional[torch.Tensor], - num_tokens_per_rdma_rank: Optional[torch.Tensor], - is_token_in_rank: torch.Tensor, - num_tokens_per_expert: Optional[torch.Tensor], - cached_num_recv_tokens: int, - cached_num_rdma_recv_tokens: int, - cached_rdma_channel_prefix_matrix: Optional[torch.Tensor], - cached_recv_rdma_rank_prefix_sum: Optional[torch.Tensor], - cached_gbl_channel_prefix_matrix: Optional[torch.Tensor], - cached_recv_gbl_rank_prefix_sum: Optional[torch.Tensor], - expert_alignment: int, - config, - ): - """High-throughput internode (NVLink + RDMA) dispatch. - - Returns ``(recv_x, recv_x_scales, recv_topk_idx, recv_topk_weights, - num_recv_tokens_per_expert_list, rdma_channel_prefix_matrix, - gbl_channel_prefix_matrix, recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, recv_gbl_channel_prefix_matrix, - recv_gbl_rank_prefix_sum, recv_src_meta, send_rdma_head, send_nvl_head)`` - to mirror the previous DeepEP-style surface.""" - assert x.dim() == 2 and x.is_contiguous() - cached_mode = cached_rdma_channel_prefix_matrix is not None - num_tokens, hidden = int(x.size(0)), int(x.size(1)) - x_element_size = x.element_size() - num_rdma_ranks = self.runtime.get_num_rdma_ranks() - num_channels = self.runtime.get_internode_dispatch_num_channels(config) - - num_topk = int(topk_idx.size(1)) if topk_idx is not None else 0 - num_scales = 0 - if x_scales is not None: - num_scales = 1 if x_scales.dim() == 1 else int(x_scales.size(1)) - - # ----- Phase A: notify (non-cached) or reuse cached layout ----- - if cached_mode: - num_recv_tokens = cached_num_recv_tokens - num_rdma_recv_tokens = cached_num_rdma_recv_tokens - rdma_channel_prefix_matrix = cached_rdma_channel_prefix_matrix - recv_rdma_rank_prefix_sum = cached_recv_rdma_rank_prefix_sum - gbl_channel_prefix_matrix = cached_gbl_channel_prefix_matrix - recv_gbl_rank_prefix_sum = cached_recv_gbl_rank_prefix_sum - num_recv_tokens_per_expert_list: List[int] = [] - num_experts = 0 - else: - assert ( - num_tokens_per_rank is not None - and num_tokens_per_rdma_rank is not None - and num_tokens_per_expert is not None - ) - num_experts = int(num_tokens_per_expert.size(0)) - num_local_experts = num_experts // self.group_size - rdma_channel_prefix_matrix = torch.empty((num_rdma_ranks, num_channels), dtype=torch.int32, device="cuda") - recv_rdma_rank_prefix_sum = torch.empty((num_rdma_ranks,), dtype=torch.int32, device="cuda") - gbl_channel_prefix_matrix = torch.empty((self.group_size, num_channels), dtype=torch.int32, device="cuda") - recv_gbl_rank_prefix_sum = torch.empty((self.group_size,), dtype=torch.int32, device="cuda") - # num_recv_tokens_per_expert and num_rdma_recv_tokens are written on the host. - num_recv_per_expert_host = torch.empty((num_local_experts,), dtype=torch.int32, device="cpu") - num_rdma_recv_host = torch.empty((1,), dtype=torch.int32, device="cpu") - num_recv_tokens = self.runtime.internode_notify_dispatch( - _ptr(rdma_channel_prefix_matrix), - _ptr(recv_rdma_rank_prefix_sum), - _ptr(gbl_channel_prefix_matrix), - _ptr(recv_gbl_rank_prefix_sum), - _ptr(num_recv_per_expert_host), - _ptr(num_rdma_recv_host), - _ptr(num_tokens_per_rank), - _ptr(num_tokens_per_rdma_rank), - _ptr(num_tokens_per_expert), - _ptr(is_token_in_rank), - num_tokens, - num_experts, - hidden, - num_scales, - num_topk, - x_element_size, - expert_alignment, - config, - _stream_ptr(), - ) - num_rdma_recv_tokens = int(num_rdma_recv_host[0].item()) - num_recv_tokens_per_expert_list = num_recv_per_expert_host.tolist() - - # ----- Phase B: allocate recv outputs (or view the recv pool) ----- - recv_x = self._alloc_internode_recv_x(num_recv_tokens, hidden, x_element_size, config, cached_mode) - recv_topk_idx = ( - torch.empty((num_recv_tokens, num_topk), dtype=torch.int64, device="cuda") if topk_idx is not None else None - ) - recv_topk_weights = ( - torch.empty((num_recv_tokens, num_topk), dtype=torch.float32, device="cuda") - if topk_weights is not None - else None - ) - recv_x_scales = ( - torch.empty((num_recv_tokens, num_scales), dtype=torch.float32, device="cuda") - if x_scales is not None - else None - ) - # The receiver-side metadata / head buffers are only produced (and only - # needed by combine) on the non-cached forward path. - if cached_mode: - recv_src_meta = None - recv_rdma_channel_prefix_matrix = None - recv_gbl_channel_prefix_matrix = None - send_rdma_head = None - send_nvl_head = None - else: - meta_bytes = self.runtime.get_source_meta_bytes() - num_max_nvl_peers = self.runtime.get_num_max_nvl_peers() - recv_src_meta = torch.empty((num_recv_tokens, meta_bytes), dtype=torch.uint8, device="cuda") - recv_rdma_channel_prefix_matrix = torch.empty( - (num_rdma_ranks, num_channels), dtype=torch.int32, device="cuda" - ) - recv_gbl_channel_prefix_matrix = torch.empty( - (self.group_size, num_channels), dtype=torch.int32, device="cuda" - ) - send_rdma_head = torch.empty((num_tokens, num_rdma_ranks), dtype=torch.int32, device="cuda") - send_nvl_head = torch.empty((num_rdma_recv_tokens, num_max_nvl_peers), dtype=torch.int32, device="cuda") - - self.runtime.internode_dispatch( - _ptr(recv_x), - _ptr(recv_x_scales), - _ptr(recv_topk_idx), - _ptr(recv_topk_weights), - _ptr(recv_src_meta), - _ptr(recv_rdma_channel_prefix_matrix), - _ptr(recv_gbl_channel_prefix_matrix), - _ptr(send_rdma_head), - _ptr(send_nvl_head), - _ptr(x), - _ptr(x_scales), - _ptr(topk_idx), - _ptr(topk_weights), - _ptr(is_token_in_rank), - _ptr(rdma_channel_prefix_matrix), - _ptr(recv_rdma_rank_prefix_sum), - _ptr(gbl_channel_prefix_matrix), - _ptr(recv_gbl_rank_prefix_sum), - num_tokens, - hidden, - num_topk, - num_scales, - num_experts, - x_element_size, - num_recv_tokens, - num_rdma_recv_tokens, - cached_mode, - config, - _stream_ptr(), - ) - return ( - recv_x, - recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rdma_channel_prefix_matrix, - gbl_channel_prefix_matrix, - recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix, - recv_gbl_rank_prefix_sum, - recv_src_meta, - send_rdma_head, - send_nvl_head, - ) - - def _alloc_internode_recv_x( - self, num_recv_tokens: int, hidden: int, x_element_size: int, config, cached_mode: bool - ) -> torch.Tensor: - """Allocate ``recv_x`` or, on the non-cached direct path, view this rank's - recv pool (so the cross-GPU forwarder writes hidden into the pool and the - direct-gather combine reads it back). The pool view is non-cached only, - matching the ``ep_use_direct`` gate in the C++ runtime.""" - if not cached_mode: - pool_ptr = self.runtime.resolve_internode_recv_x_buffer(num_recv_tokens, hidden, x_element_size, config) - if pool_ptr != 0: - return _bf16_view(pool_ptr, num_recv_tokens, hidden, owner=self) - return torch.empty((num_recv_tokens, hidden), dtype=torch.bfloat16, device="cuda") - - def internode_combine( - self, - x: torch.Tensor, - topk_weights: Optional[torch.Tensor], - src_meta: torch.Tensor, - is_combined_token_in_rank: torch.Tensor, - rdma_channel_prefix_matrix: torch.Tensor, - rdma_rank_prefix_sum: torch.Tensor, - gbl_channel_prefix_matrix: torch.Tensor, - combined_rdma_head: torch.Tensor, - combined_nvl_head: torch.Tensor, - config, - ): - """Returns ``(combined_x, combined_topk_weights|None)``.""" - assert x.dim() == 2 and x.is_contiguous() - num_tokens, hidden = int(x.size(0)), int(x.size(1)) - num_combined_tokens = int(is_combined_token_in_rank.size(0)) - num_topk = int(topk_weights.size(1)) if topk_weights is not None else 0 - - combined_x = torch.empty((num_combined_tokens, hidden), dtype=torch.bfloat16, device="cuda") - combined_topk_weights = ( - torch.empty((num_combined_tokens, num_topk), dtype=torch.float32, device="cuda") - if topk_weights is not None - else None - ) - self.runtime.internode_combine( - _ptr(combined_x), - _ptr(combined_topk_weights), - _ptr(x), - _ptr(topk_weights), - _ptr(src_meta), - _ptr(is_combined_token_in_rank), - _ptr(rdma_channel_prefix_matrix), - _ptr(rdma_rank_prefix_sum), - _ptr(gbl_channel_prefix_matrix), - _ptr(combined_rdma_head), - _ptr(combined_nvl_head), - num_tokens, - num_combined_tokens, - hidden, - num_topk, - x.element_size(), - config, _stream_ptr(), ) return combined_x, combined_topk_weights @@ -627,19 +308,12 @@ class HighThroughputBackend: self.num_sms, config.nvl_chunked_send, config.nvl_chunked_recv, - config.rdma_chunked_send, - config.rdma_chunked_recv, ) - hidden_bytes = self.hidden_size * torch.tensor([], dtype=torch.bfloat16).element_size() - num_nvl_bytes = self._cfg.get_nvl_buffer_size_hint(hidden_bytes, self.world_size) - num_rdma_bytes = self._cfg.get_rdma_buffer_size_hint(hidden_bytes, self.world_size) - self._is_internode = num_rdma_bytes > 0 + hidden_bytes = self.hidden_size * torch.empty((), dtype=torch.bfloat16).element_size() self._runtime = HighThroughputRuntime( comm, - num_nvl_bytes=num_nvl_bytes, - num_rdma_bytes=num_rdma_bytes, - low_latency_mode=False, - num_qps_per_rank=config.num_rdma_qps_per_rank, + max_hidden_bytes=hidden_bytes, + config=self._cfg, ) def is_available(self) -> bool: @@ -649,7 +323,7 @@ class HighThroughputBackend: return self._runtime.is_internode_available() def is_internode(self) -> bool: - return self._is_internode + return self._runtime.is_internode_available() def dispatch( self, @@ -683,72 +357,16 @@ class HighThroughputBackend: cache = getattr(previous_handle, "_dispatch_cache", None) if previous_handle is not None else None if cache is not None: num_tokens_per_rank = cache["num_tokens_per_rank"] - num_tokens_per_rdma_rank = cache["num_tokens_per_rdma_rank"] num_tokens_per_expert = cache["num_tokens_per_expert"] is_token_in_rank = cache["is_token_in_rank"] else: ( num_tokens_per_rank, - num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, ) = self._runtime.get_dispatch_layout(topk_ids, self.num_experts) - if self._is_internode: - ( - recv_x, - _recv_x_scales, - _recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - _rdma_channel_prefix_matrix, - _gbl_channel_prefix_matrix, - recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix, - _recv_gbl_rank_prefix_sum, - recv_src_meta, - send_rdma_head, - send_nvl_head, - ) = self._runtime.internode_dispatch( - input, - None, - topk_ids, - weights, - num_tokens_per_rank, - num_tokens_per_rdma_rank, - is_token_in_rank, - num_tokens_per_expert, - 0, - 0, - None, - None, - None, - None, - self.expert_alignment, - self._cfg, - ) - combine_context = RowMajorInternodeCombineContext( - recv_topk_weights=recv_topk_weights, - src_meta=recv_src_meta, - is_token_in_rank=is_token_in_rank, - recv_rdma_channel_prefix_matrix=recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum=recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix=recv_gbl_channel_prefix_matrix, - send_rdma_head=send_rdma_head, - send_nvl_head=send_nvl_head, - ) - dispatch_cache = ( - cache - if cache is not None - else { - "num_tokens_per_rank": num_tokens_per_rank, - "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, - "num_tokens_per_expert": num_tokens_per_expert, - "is_token_in_rank": is_token_in_rank, - } - ) - elif cache is not None: + if cache is not None: ( recv_x, _recv_x_scales, @@ -760,7 +378,7 @@ class HighThroughputBackend: recv_channel_prefix_matrix, recv_src_idx, send_head, - ) = self._runtime.intranode_dispatch( + ) = self._runtime.dispatch( input, None, None, @@ -772,9 +390,8 @@ class HighThroughputBackend: cache["rank_prefix_matrix"], cache["channel_prefix_matrix"], self.expert_alignment, - self._cfg, ) - combine_context = RowMajorIntranodeCombineContext( + combine_context = RowMajorCombineContext( recv_topk_weights=recv_topk_weights, src_idx=recv_src_idx, rank_prefix_matrix=rank_prefix_matrix, @@ -794,7 +411,7 @@ class HighThroughputBackend: recv_channel_prefix_matrix, recv_src_idx, send_head, - ) = self._runtime.intranode_dispatch( + ) = self._runtime.dispatch( input, None, topk_ids, @@ -806,9 +423,8 @@ class HighThroughputBackend: None, None, self.expert_alignment, - self._cfg, ) - combine_context = RowMajorIntranodeCombineContext( + combine_context = RowMajorCombineContext( recv_topk_weights=recv_topk_weights, src_idx=recv_src_idx, rank_prefix_matrix=rank_prefix_matrix, @@ -817,7 +433,6 @@ class HighThroughputBackend: ) dispatch_cache = { "num_tokens_per_rank": num_tokens_per_rank, - "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, "num_tokens_per_expert": num_tokens_per_expert, "is_token_in_rank": is_token_in_rank, "rank_prefix_matrix": rank_prefix_matrix, @@ -838,12 +453,7 @@ class HighThroughputBackend: quant=output_info.quant, layout=output_info.layout, ) - handle_cls = ( - RowMajorInternodeDispatchHandle - if isinstance(combine_context, RowMajorInternodeCombineContext) - else RowMajorIntranodeDispatchHandle - ) - handle = handle_cls(output_info=output_info, combine_context=combine_context) + handle = RowMajorDispatchHandle(output_info=output_info, combine_context=combine_context) # The torch-free HT runtime orders its work on the caller's CUDA stream # (no separate event handle), so there is nothing to attach here. handle._event = None # type: ignore[attr-defined] @@ -867,33 +477,15 @@ class HighThroughputBackend: self, expert_output: torch.Tensor, handle: DispatchHandle, out: Optional[torch.Tensor] ) -> torch.Tensor: self._validate_combine_inputs(expert_output, handle) - if isinstance(handle, RowMajorInternodeDispatchHandle): - context = handle.combine_context - combined_x, _combined_w = self._runtime.internode_combine( - expert_output, - context.recv_topk_weights, - context.src_meta, - context.is_token_in_rank, - context.recv_rdma_channel_prefix_matrix, - context.recv_rdma_rank_prefix_sum, - context.recv_gbl_channel_prefix_matrix, - context.send_rdma_head, - context.send_nvl_head, - self._cfg, - ) - elif isinstance(handle, RowMajorIntranodeDispatchHandle): - context = handle.combine_context - combined_x, _combined_w = self._runtime.intranode_combine( - expert_output, - context.recv_topk_weights, - context.src_idx, - context.rank_prefix_matrix, - context.recv_channel_prefix_matrix, - context.send_head, - self._cfg, - ) - else: - raise ValueError("DispatchHandle does not contain row-major combine context") + context = handle.combine_context + combined_x, _combined_w = self._runtime.combine( + expert_output, + context.recv_topk_weights, + context.src_idx, + context.rank_prefix_matrix, + context.recv_channel_prefix_matrix, + context.send_head, + ) if out is not None: out.copy_(combined_x) return out @@ -925,11 +517,9 @@ class HighThroughputBackend: raise ValueError("weights shape must match topk_ids") def _validate_combine_inputs(self, expert_output, handle) -> None: - if not isinstance(handle, (RowMajorIntranodeDispatchHandle, RowMajorInternodeDispatchHandle)): + if not isinstance(handle, RowMajorDispatchHandle): raise TypeError("handle must be a DispatchHandle returned by dispatch") if expert_output.dim() != 2 or not expert_output.is_contiguous(): raise ValueError("expert_output must be a contiguous [total_recv_tokens, hidden] tensor") if expert_output.size(1) != self.hidden_size: raise ValueError(f"expert_output hidden size {expert_output.size(1)} != configured {self.hidden_size}") - if self._is_internode != isinstance(handle, RowMajorInternodeDispatchHandle): - raise ValueError("handle transport does not match this communicator") diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index 451a984b..35fde230 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -56,8 +56,7 @@ class MoECommunicatorConfig: # Quantization defaults quant: Optional[QuantConfig] = None - # Transport / launch tuning - num_rdma_qps_per_rank: int = 12 + # Launch tuning num_sms: int = 20 low_latency_num_blocks: int = 130 low_latency_combine_mode: CombineMode = CombineMode.RANK_LOCAL_REDUCE @@ -67,8 +66,6 @@ class MoECommunicatorConfig: expert_alignment: int = 1 nvl_chunked_send: int = 8 nvl_chunked_recv: int = 256 - rdma_chunked_send: int = 16 - rdma_chunked_recv: int = 128 # MLP-facing dispatch output. @@ -118,8 +115,8 @@ class ExpertMajorCombineContext: @dataclass -class RowMajorIntranodeCombineContext: - """Combine context for row-major intranode dispatch output.""" +class RowMajorCombineContext: + """Combine context for row-major high-throughput dispatch output.""" recv_topk_weights: Optional[torch.Tensor] src_idx: torch.Tensor @@ -128,21 +125,7 @@ class RowMajorIntranodeCombineContext: send_head: torch.Tensor -@dataclass -class RowMajorInternodeCombineContext: - """Combine context for row-major internode dispatch output.""" - - recv_topk_weights: Optional[torch.Tensor] - src_meta: torch.Tensor - is_token_in_rank: torch.Tensor - recv_rdma_channel_prefix_matrix: torch.Tensor - recv_rdma_rank_prefix_sum: torch.Tensor - recv_gbl_channel_prefix_matrix: torch.Tensor - send_rdma_head: torch.Tensor - send_nvl_head: torch.Tensor - - -CombineContext = Union[ExpertMajorCombineContext, RowMajorIntranodeCombineContext, RowMajorInternodeCombineContext] +CombineContext = Union[ExpertMajorCombineContext, RowMajorCombineContext] # Opaque dispatch handles returned by dispatch() and consumed by combine(). @@ -161,13 +144,8 @@ class ExpertMajorDispatchHandle(DispatchHandle): @dataclass -class RowMajorIntranodeDispatchHandle(DispatchHandle): - combine_context: RowMajorIntranodeCombineContext - - -@dataclass -class RowMajorInternodeDispatchHandle(DispatchHandle): - combine_context: RowMajorInternodeCombineContext +class RowMajorDispatchHandle(DispatchHandle): + combine_context: RowMajorCombineContext # Optional async/overlap configuration. diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index ae775d05..f58cc276 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -7,7 +7,7 @@ # Two backends share one module, both with a torch-free, raw-pointer (uintptr_t) # nanobind API so the module never links libtorch: # - Low-latency (LL): moe_runtime.cc + low_latency/{dispatch,combine}.cu. -# - High-throughput (HT): ht_runtime.cc + ht/kernels/*.cu, a DeepEP-style +# - High-throughput (HT): ht_runtime.cc + ht/*.cu, a DeepEP-style # runtime de-torched to the same pointer boundary as the LL runtime (dynamic # recv sizing via the two-phase notify -> allocate -> dispatch API). @@ -44,15 +44,15 @@ set(EP_SOURCES low_latency/combine.cu # High-throughput (DeepEP-style) backend (torch-free, raw-pointer API). ht_runtime.cc - ht/kernels/internode.cu - ht/kernels/intranode_kernel.cu - ht/kernels/runtime.cu + ht/layout.cu + ht/intranode_kernel.cu + ht/runtime.cu ) if(MSCCLPP_USE_ROCM) # ROCm port of the EP kernels is not supported yet; see src/ext/ep/README.md. message(WARNING "mscclpp_ep: ROCm build path not implemented, falling back to CXX compile.") - file(GLOB_RECURSE EP_CU_SOURCES low_latency/*.cu ht/kernels/*.cu) + file(GLOB_RECURSE EP_CU_SOURCES low_latency/*.cu ht/*.cu) set_source_files_properties(${EP_CU_SOURCES} PROPERTIES LANGUAGE CXX) endif() @@ -73,10 +73,6 @@ target_include_directories(mscclpp_ep_cpp PRIVATE ) target_link_libraries(mscclpp_ep_cpp PRIVATE mscclpp ${GPU_LIBRARIES} Threads::Threads) -# The HT NCCL-EP dispatch/combine (sender-direct + flat all-sender + TMA) path -# is compiled in unconditionally for this branch (the validated GB200 fast path). -target_compile_definitions(mscclpp_ep_cpp PRIVATE EP_DISPATCH_NCCLEP) - # The EP CUDA kernels call constexpr __host__ helpers (std::max / std::pair) and # use extended lambdas in device code. These flags were previously injected # implicitly by Torch's CMake package; add them explicitly now that the module @@ -88,14 +84,6 @@ if(MSCCLPP_USE_CUDA) ) endif() -# Number of intra-node NVLink peers compiled into the EP kernels. -# - 8 (default) matches DeepEP upstream and H100/HGX 8-GPU nodes. -# - 4 is required for Azure GB200 NVL72 (4 GPUs per NUMA host). -set(MSCCLPP_EP_NUM_MAX_NVL_PEERS "8" CACHE STRING - "Compile-time NUM_MAX_NVL_PEERS for the EP kernels (8 for HGX, 4 for GB200)") -target_compile_definitions(mscclpp_ep_cpp PRIVATE - NUM_MAX_NVL_PEERS=${MSCCLPP_EP_NUM_MAX_NVL_PEERS}) - # Kernel-side debug timeout (~10s) — set via: # -DMSCCLPP_EP_KERNEL_DEBUG_TIMEOUT=ON option(MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 3a4fd1b9..4a9f75de 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -1,567 +1,146 @@ # MSCCL++ Expert-Parallel (EP) extension -A torch-free nanobind extension for MoE `dispatch` / `combine` primitives in -MSCCL++. The module builds two active backends: +The EP extension is a torch-free nanobind module for MoE dispatch and combine. +It builds two backends: -- **Low-latency (LL)**: `MoERuntime` / `MoECommunicator(mode=LOW_LATENCY)`, - backed by `low_latency/dispatch.cu` and `low_latency/combine.cu`. -- **High-throughput (HT)**: `ExpertParallelRuntime` / - `MoECommunicator(mode=HIGH_THROUGHPUT)`, backed by `ht_runtime.cc` and - `ht/kernels/*`. +- **Low latency (LL)**: `MoERuntime`, backed by + `low_latency/{dispatch,combine}.cu`. +- **High throughput (HT)**: `ExpertParallelRuntime`, backed by + `ht_runtime.cc` and the CUDA sources under `ht/`. ## Status -| Feature | Status | -|----------------------------------|--------| -| `mscclpp_ep_cpp` module | ✅ builds LL + HT backends when `MSCCLPP_BUILD_EXT_EP=ON` | -| LL dispatch/combine | ✅ validated on 8-rank H100; cross-node requires one detected GPU IPC fabric domain | -| HT dispatch/combine | ✅ active DeepEP-style backend with intranode and internode paths | -| HT GB200 direct/flat fast paths | ✅ runtime-gated by `MSCCLPP_EP_DIRECT`, `MSCCLPP_EP_INTRA_DIRECT`, and `MSCCLPP_EP_FLAT` | -| LL topology discovery | ✅ automatic node + NVML fabric-domain detection through `Bootstrap` | -| Python frontend `mscclpp.ep` | ✅ `MoECommunicator` selects LL or HT by `MoEMode` | +| Feature | Status | +|---|---| +| LL dispatch/combine | Validated on Hopper and newer GPUs | +| HT dispatch/combine | Supports 2, 4, or 8 ranks in one GPU IPC/NVL fabric domain | +| HT RDMA/IB fallback | Not supported | +| Python frontend | `mscclpp.ep.MoECommunicator` selects LL or HT with `MoEMode` | +| ROCm | Not supported | -### LL topology and transport +## Runtime architecture -The optimized LL backend uses direct peer mappings: +### Low latency -- Same-host peers use regular CUDA IPC mappings. -- Cross-host peers use CUDA fabric handles when all ranks belong to one - NVML-reported GPU fabric domain, such as GB200 NVL72 with `nvidia-imex`. -- LL always allocates its symmetric buffer with CUDA physical memory - (`cuMemCreate`/`cuMemMap`); there is no `cudaMalloc` fallback. -- `BaseMemoryChannel` handles are used only for signal/wait synchronization; - LL payload data does not flow through `MemoryChannel` or `PortChannel`. -- There is no CPU-proxy/IB fallback in the optimized LL backend. If the full - communicator is not one GPU IPC domain, `MoERuntime::isAvailable()` is false. +LL allocates CUDA physical symmetric memory and maps peer buffers through the +existing `mscclpp::Communicator`. Payloads use direct peer mappings; +`BaseMemoryChannel` handles are used only for synchronization. -Topology is automatic. `TcpBootstrap` detects: +The optimized LL backend is available when all participating ranks belong to +one detected GPU IPC domain. That domain may span hosts when CUDA fabric handles +and the required fabric services are available. -1. ranks per host from bootstrap peer addresses; -2. ranks per GPU IPC domain from NVML fabric `clusterUuid + cliqueId`; -3. host-local IPC domains as a safe fallback when NVML fabric information is - unavailable or incomplete. +### High throughput -LL therefore does **not** require topology environment variables such as -`MSCCLPP_EP_LOCAL_WORLD_SIZE`, `MSCCLPP_EP_FABRIC_IPC`, an NVML-domain ID, or -`NCCL_MNNVL_ENABLE`. Socket/HCA variables are only optional launch overrides -when the generic bootstrap or an HT/IB test selects the wrong interface. +HT follows the same direct-mapping resource model: -### HT / legacy proxy sharding +1. Python passes the existing `mscclpp::Communicator` into + `ExpertParallelRuntime`. +2. Each rank allocates a ring/FIFO region plus a CUDA physical direct receive + pool. Same-host ring buffers use `cudaMalloc` runtime IPC for peak latency; + a fabric-domain job spanning hosts uses CUDA physical memory for the ring. +3. The runtime exchanges and maps those allocations with + `Communicator::sendMemory` / `recvMemory`. +4. Dispatch and combine launch directly on the caller's CUDA stream. -A single `mscclpp::ProxyService` is one CPU host thread driving one -FIFO. With 8 GPUs / node sharing one proxy, the host thread becomes the -bottleneck for cross-node proxy traffic. The HT runtime therefore allocates `N` -ProxyServices and shards `PortChannel`s across them by `(qp_idx, -dst_rank)`. +The detected GPU IPC domain may span multiple hosts, such as an NVL fabric +domain with CUDA fabric handles. HT does not create a private bootstrap, proxy +service, RDMA channel, NVLS multicast object, or private communication stream, +and it has no RDMA/IB fallback outside that domain. -- Default: `8` on Hopper (sm_90), `1` on Blackwell / sm_100+ (NVSwitch). -- Override at runtime: `MSCCLPP_EP_NUM_PROXIES=` (clamped to ≥1). - Rank 0 prints the resolved value at construction. -- Sweet spot on 2×H100×8 is `N=8`; `N=12` over-subscribes the host CPUs - and collapses throughput. +The HT dispatch API remains two-phase because the receive token count is data +dependent: -### Known limitations +1. `notify_dispatch` exchanges counts and produces prefix matrices. +2. Python allocates the exact receive tensors. +3. `dispatch` moves token data and metadata. -- ROCm is not supported for the EP extension yet. -- LL requires Hopper or newer and CUDA physical-memory allocation support. -- LL currently supports BF16 input, optional FP8 E4M3 dispatch output, and - `DispatchLayout.EXPERT_MAJOR`. -- HT currently supports BF16 input and `DispatchLayout.FLAT`; quantized HT - dispatch is not implemented. -- Optimized LL requires all participating ranks to share one detected CUDA IPC - domain; ordinary multi-node H100 jobs should use HT. -- HT direct and flat paths are GB200/NVL72 optimizations. Leave the env vars - unset for the baseline DeepEP-style HT path. +Cached dispatch reuses the previous receive count and prefix matrices. + +## HT data paths + +The baseline path uses the DeepEP-style intranode ring. Optional runtime paths +use the peer-mapped receive pool: + +- `MSCCLPP_EP_INTRA_DIRECT=1`: send hidden rows directly to their final receive + slots. The physical receive pool is allocated and mapped only when this flag + is enabled before runtime construction. +- `MSCCLPP_EP_INTRA_ALLSENDER=0|1`: controls the all-sender dispatch path when + direct dispatch and TMA combine are enabled; default is enabled. +- `MSCCLPP_EP_COMBINE_TMA=0|1`: selects TMA direct-gather combine; default is + enabled when its inputs are available. +- `MSCCLPP_EP_DISPATCH_NSM=`: overrides the dispatch block count, rounded + down to an even value. +- `MSCCLPP_EP_COMBINE_NSM=`: overrides the TMA combine block count. + +The persistent HT configuration contains only: + +| Field | Meaning | +|---|---| +| `num_sms` | Maximum HT communication block budget | +| `num_max_nvl_chunked_send_tokens` | Ring send chunk size | +| `num_max_nvl_chunked_recv_tokens` | Ring receive capacity | ## Build -Python installs build the EP extension by default: +Python builds include the EP extension by default: ```bash -python -m pip install . -# Optional CuPy dependency: -python -m pip install ".[cuda12]" +python3 -m pip install . ``` -The EP extension targets CUDA architectures **90 or newer**. Plain CMake builds -can enable it explicitly: +Plain CMake builds can enable it explicitly: ```bash -cmake -S . -B build \ - -DMSCCLPP_BUILD_EXT_EP=ON -cmake --build build -j -``` - -This produces `mscclpp_ep_cpp.so` — a nanobind extension module. -The Python frontend picks it up automatically: - -```python -import mscclpp.ep as ep -moe_comm = ep.MoECommunicator(...) -``` - -### Build-time CMake options - -| Variable | Default | Meaning | -|---------------------------------------|---------|---------------------------------------------------------------| -| `MSCCLPP_BUILD_EXT_EP` | `ON` in Python wheels | Build the EP extension at all | -| `MSCCLPP_EP_NUM_MAX_NVL_PEERS` | `8` | HT compile-time NVLink peer capacity; optimized LL topology is detected at runtime | -| `MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT` | `OFF` | Use a short ~10s kernel spin timeout (default is ~100s) | - -### Azure GB200 (NVL72, 4 GPUs / NUMA host) - -GB200 NVL72 nodes expose **4 GPUs per NUMA host** (not 8 like HGX -H100). The optimized LL backend discovers this topology from bootstrap/NVML -and needs no GB200-specific build option. HT kernels still use the compile-time -NVLink peer capacity and should be built with: - -```bash -cmake -S . -B build \ - -DMSCCLPP_BUILD_EXT_EP=ON \ - -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=4 +cmake -S . -B build -DMSCCLPP_BUILD_EXT_EP=ON cmake --build build -j 64 - -python3 -m pip install . \ - --config-settings=cmake.define.MSCCLPP_EP_NUM_MAX_NVL_PEERS=4 ``` -Add `-DMSCCLPP_EP_KERNEL_DEBUG_TIMEOUT=ON` only when triaging hangs (it -shortens the kernel-side spin timeout from ~100s to ~10s). +The EP extension requires CUDA architecture 90 or newer. -Runtime prerequisites on GB200: +Available CMake options: -- CUDA Toolkit ≥ 12.5 (the `cuCtxCreate` proxy-context path uses the - 4-arg signature added in 12.5; older toolkits compile against the - 3-arg fallback automatically). -- Driver ≥ 555 with nvidia-imex configured so cuMem fabric handles - (`POSIX_FD | FABRIC`) can be exchanged across nodes. -- NVML must report a completed fabric state with a common cluster UUID and - clique ID for all participating ranks. Bootstrap uses this information to - form the IPC domain automatically. -- The symmetric allocation requests POSIX-FD handles and fabric handles when - the device exposes them. If bootstrap does not detect one IPC domain spanning - every rank, LL reports unavailable rather than attempting an IB fallback. -- RT priority is required by NCCL/glibc. On each node: +| Variable | Default | Meaning | +|---|---:|---| +| `MSCCLPP_BUILD_EXT_EP` | Python builds: `ON` | Build the EP extension | +| `MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT` | `OFF` | Use a shorter kernel spin timeout | - ```bash - sudo tee -a /etc/security/limits.conf > /dev/null <<'EOF' - * soft rtprio 99 - * hard rtprio 99 - EOF - # Re-login so `ulimit -r` reports 99. - ``` -- `nvidia-imex` must be active on every node with an identical - `/etc/nvidia-imex/nodes_config.cfg` listing all node IPs. Verify: +## Source layout - ```bash - sudo systemctl status nvidia-imex - sudo cat /etc/nvidia-imex/nodes_config.cfg - ls /dev/nvidia-caps-imex-channels/ # channel0 must exist - ``` - -No LL-specific runtime environment variables are required. Optional overrides -are only needed when generic launch components choose the wrong interface: - -```bash -export NCCL_SOCKET_IFNAME=enP22p1s0f1 -export MSCCLPP_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME -export GLOO_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME -# HT/IB only, if automatic HCA selection chooses a down/bond device: -export MSCCLPP_HCA_DEVICES=mlx5_0,mlx5_1,mlx5_2,mlx5_3 -``` - -`NCCL_IB_DISABLE` and `NCCL_MNNVL_ENABLE` configure NCCL, not the MSCCL++ LL -runtime. Set them only when the surrounding PyTorch/NCCL test launcher needs -those overrides. - -HT runtime knobs (env vars exposed by `test_intranode_multirank.py` / -`test_internode_multirank.py` — defaults below are the test-script defaults, -**not** the `ep.Config(...)` constructor defaults): - -| Variable | Maps to (`ep.Config` field) | Default | Notes | -|---------------------------|--------------------------------------|--------:|-----------------------------------------| -| `MSCCLPP_EP_NUM_SMS` | `num_sms` | `20` | HT intranode test only. Try `64` on GB200 intranode for `dispatch` BW. | -| `MSCCLPP_EP_NSM` | `num_sms` | `152` | HT internode test only. | -| `MSCCLPP_EP_NVL_SEND` | `num_max_nvl_chunked_send_tokens` | `8` | Must be `<` `MSCCLPP_EP_NVL_RECV`. | -| `MSCCLPP_EP_NVL_RECV` | `num_max_nvl_chunked_recv_tokens` | `256` | Scales NVL ring buffer linearly. | -| `MSCCLPP_EP_RDMA_SEND` | `num_max_rdma_chunked_send_tokens` | `16` | Internode only. | -| `MSCCLPP_EP_RDMA_RECV` | `num_max_rdma_chunked_recv_tokens` | `128` | Scale **down** as `num_rdma_ranks` grows (4n→128, 8n→64, 16n→32) to keep the RDMA buffer under the 2 GiB `INT_MAX` limit. | -| `MSCCLPP_EP_DIRECT` | — (runtime `getenv`) | unset | **GB200 / NVL72 internode only.** `1` enables the sender direct-write dispatch **and** receiver gather-direct combine (see [GB200 direct-path optimization](#gb200-direct-path-optimization-mscclpp_ep_direct--mscclpp_ep_intra_direct)). Unset = byte-identical 2-hop baseline. | -| `MSCCLPP_EP_INTRA_DIRECT` | — (runtime `getenv`) | unset | **GB200 single-node only.** `1` enables sender direct-write for the intra-node kernel. Independent of `MSCCLPP_EP_DIRECT` (see below). | -| `MSCCLPP_EP_FLAT` | — (runtime `getenv`) | unset | **GB200 / NVL72 internode only; requires `MSCCLPP_EP_DIRECT=1`.** `1` enables the flat all-sender dispatch (removes the forwarder / coordinator / receiver roles; per-token metadata goes straight to the dest recv pool) + direct-gather combine (see [flat all-sender path](#gb200-flat-all-sender-path-mscclpp_ep_flat)). Unset = the `MSCCLPP_EP_DIRECT` path. | -| `MSCCLPP_EP_DISPATCH_NSM` | — (runtime `getenv`) | `num_sms`/2 | **Flat path only (`MSCCLPP_EP_FLAT=1`).** Sets the all-sender dispatch block count independently of `num_sms`; clamped to `[1, num_sms]` (the flat path has no forwarder, so it can use the full SM budget — e.g. `num_sms=16` + `DISPATCH_NSM=16` → 16 blocks). The RDMA/NVL buffers are auto-sized to match. Lower it to free SMs for overlapping compute, or set it to `num_sms` to maximize dispatch throughput. | -| `MSCCLPP_EP_COMBINE_NSM` | — (runtime `getenv`) | `num_sms` | **Flat path only (`MSCCLPP_EP_FLAT=1`).** Caps the combine block count independently of `num_sms`; clamped to `[2, num_sms]`. Flat combine saturates ~76 blocks. | -| `MSCCLPP_EP_COMBINE_TMA` | — (runtime `getenv`) | `1` (on) | **Flat path only (`MSCCLPP_EP_FLAT=1` + `MSCCLPP_EP_DIRECT=1`).** Selects the combine flat-gather implementation. Default (`1`) uses the TMA-staged gather (`cp.async.bulk` contributor rows → SMEM → reduce), which hides remote-NVLink read latency via the async copy engine and wins at every channel count and node scale. `0` falls back to the synchronous register-MLP gather (lean kernel ≤14 channels, else the unified flat branch). | - -Validated 16-node (64-rank) configs on Azure GB200 NVL72 (HIDDEN=7168, -tokens=4096, experts=256, topk=8): - -- HT internode: `NVL_SEND=8 NVL_RECV=256 RDMA_SEND=8 RDMA_RECV=32` → - dispatch ~**2 006 GB/s** agg, combine ~**2 011 GB/s** agg. - -### GB200 direct-path optimization (`MSCCLPP_EP_DIRECT` / `MSCCLPP_EP_INTRA_DIRECT`) - -On GB200 NVL72 the cross-node "RDMA send" is a cuMem-fabric VA write over -NVLink, so a sender can reach any peer GPU's recv pool directly. Two -env-gated flags exploit this to remove the classic DeepEP 2-hop (forwarder -transpose on dispatch / ring drain on intra-node), turning a -structural-/handshake-bound copy into a bandwidth-bound 1-hop that scales -with SM count. Both default **off**, and the binary is byte-identical to the -2-hop baseline when unset (one `.so`, runtime `getenv`-gated). - -- **`MSCCLPP_EP_DIRECT=1`** — internode HT (`test_internode_multirank.py`). - One master flag that turns on **both**: - - *Dispatch sender direct-write:* `kRDMASender` writes each token's hidden - straight to `recv_pool_global_ptrs[dst] + header + idx*hidden_bytes` - (1 hop), instead of the forwarder transpose to local NVL peers (2 hops). - - *Combine receiver gather-direct:* each token gathers its top-k expert - contributions directly from the peer recv pools and reduces locally, - skipping the `nvl_channel + forwarder + rdma_channel` path. - - *Prerequisite:* cross-node fabric-IPC pool mapping — auto-detected on - GB200; force with `MSCCLPP_EP_FABRIC_IPC=1` if needed. The combine input - must live in the recv pool (DeepEP contract; satisfied by the round-trip, - where combine input = dispatch output). -- **`MSCCLPP_EP_INTRA_DIRECT=1`** — single-node - (`test_intranode_multirank.py`), a **separate** flag for the intra-node - kernel. The sender writes hidden into the destination GPU's peer-mapped - pool and the receiver skips the ring drain. - -Measured on Azure GB200 NVL72 (2 nodes × 4 GPU, HIDDEN=7168, tokens=4096, -topk=8, experts=256): - -- HT internode dispatch+combine round-trip (NSM=20): **~3 850 µs → ~1 750 µs - (−54 %)** vs the 2-hop baseline; the win grows with node count. -- HT intra-node `INTRA_DIRECT=1` dispatch scales **980 µs (16 SM) → 285 µs - (152 SM, ~206 GB/s per rank)**, whereas the 2-hop baseline is SM-flat at - ~3.8 ms. - -> **Single-node launch note.** Launch the intra-node test with an explicit -> `127.0.0.1` rendezvous (not `torchrun --standalone`, whose hostname -> rendezvous is not DNS-resolvable on these nodes) and set -> `NCCL_NET_PLUGIN=none` + `NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3` so NCCL's -> built-in IB probe does not crash `ep.MoECommunicator` construction. - -### GB200 flat all-sender path (`MSCCLPP_EP_FLAT`) - -`MSCCLPP_EP_FLAT=1` (requires `MSCCLPP_EP_DIRECT=1`) takes the direct path one -step further. The `MSCCLPP_EP_DIRECT` dispatch still runs a sender + forwarder -pair per channel and routes per-token metadata through the 2-hop RDMA ring → -forwarder → NVL-receiver pipeline; the flat path removes the forwarder, -sender-coordinator, and NVL-receiver roles entirely, so dispatch launches -**all-sender** (one block per channel) and each sender writes its tokens' -metadata straight into the destination recv pool. Combine is a flat -direct-gather (no forwarder); by default it uses a **TMA-staged** gather that -streams each contributor's hidden chunks through shared memory via `cp.async.bulk` -(`MSCCLPP_EP_COMBINE_TMA=1`, the default), letting the async copy engine hide the -remote-NVLink read latency. A small post-dispatch drain copies the pool metadata -into the `recv_*` output tensors. - -Because both legs are now bandwidth-bound 1-hops, they saturate the NVLink -write/read ceiling well below the full SM grid, so the dispatch and combine -block counts can be capped **independently** of `num_sms` to free SMs for -overlapping compute: - -- `MSCCLPP_EP_DISPATCH_NSM=` — flat dispatch block count, clamped to - `[1, num_sms]` (the flat path has no forwarder, so it can use the full SM - budget; the RDMA/NVL buffers are auto-sized to the chosen count). - Default = `num_sms/2`. -- `MSCCLPP_EP_COMBINE_NSM=` — flat combine block count, clamped to - `[2, num_sms]`. Default = `num_sms`. - -Both are flat-only and leave the `MSCCLPP_EP_DIRECT` / 2-hop paths byte- -identical when unset. Measured on Azure GB200 NVL72 (2 nodes × 4 GPU, -HIDDEN=7168, tokens=4096, topk=8, experts=256, `MSCCLPP_EP_NSM=152`): - -- Dispatch (combine grid fixed): `DISPATCH_NSM` 16→**698 µs**, 32→544, - 64→498, 76→**495 µs** — knee at ~64 blocks. -- Combine (dispatch grid fixed): `COMBINE_NSM` 16→**1006 µs**, 32→577, - 64→514, 152→**451 µs** — knee at ~76 blocks. - -So flat dispatch reaches its floor at ~64 of 76 possible blocks and combine at -~76 of 152, leaving the remaining SMs free for the model's compute to overlap -the communication. - -## Layout - -``` +```text src/ext/ep/ -├── CMakeLists.txt — builds mscclpp_ep_cpp (nanobind) -├── README.md — this file -├── moe_runtime.hpp / .cc — LL MoE runtime state and raw-pointer dispatch/combine -├── config.hpp — LL layout helpers and size hints -├── bindings.cpp — nanobind module definition -├── ht_runtime.hpp / .cc — HT runtime state and raw-pointer dispatch/combine -├── ht/ — active HT kernel/config sources -│ ├── buffer.hpp / buffer.cc +├── bindings.cpp +├── moe_runtime.{cc,hpp} +├── ht_runtime.{cc,hpp} +├── ht/ │ ├── config.hpp -│ ├── event.hpp -│ └── kernels/ -│ ├── buffer.cuh -│ ├── runtime.cu -│ ├── intranode_kernel.cu -│ ├── internode.cu -│ └── internode_ncclep.cuh -├── include/ — shared EP headers and quantization helpers +│ ├── buffer.cuh +│ ├── layout.cu +│ ├── intranode_kernel.cu +│ └── runtime.cu +├── include/ └── low_latency/ - ├── dispatch.cu — optimized LL dispatch - ├── combine.cu — optimized LL combine - └── config.cuh — LL launch/workspace configuration - -python/mscclpp/ep/ -├── __init__.py — public exports -├── communicator.py — backend-selecting public API -├── low_latency.py — LL torch.Tensor frontend -└── high_throughput.py — HT torch.Tensor frontend - -test/python/ep/ -├── test_intranode_multirank.py — intranode HT dispatch+combine -├── test_internode_multirank.py — internode HT dispatch+combine -├── test_low_latency_multirank.py — LL correctness + CUDA Graph -├── ep_bench_ll.py — Python NCCL-style LL benchmark -├── mscclpp_ep_bench.cu — pure-C++ NCCL-style LL benchmark -└── run_ep_bench.py — unified MSCCL++ / NCCL-EP driver + ├── config.cuh + ├── dispatch.cu + └── combine.cu ``` -## Running the tests +## Validation -### Test prerequisites - -The Python tests are launched through `torchrun` / `mpirun` and require -PyTorch + a few support packages in the active environment. A minimal -install (matches the GB200 reference setup): +Build the extension, then run the single-node HT test: ```bash -# Conda env (any Python >= 3.10). Use the appropriate Miniconda variant -# for the host arch (`aarch64` shown; use `x86_64` on x86 clusters). -wget -O /tmp/Miniconda3-latest-Linux-aarch64.sh \ - https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh -bash /tmp/Miniconda3-latest-Linux-aarch64.sh -b -p $HOME/miniconda3 -source $HOME/miniconda3/etc/profile.d/conda.sh -conda create -n torch python=3.14 -y -conda activate torch - -# Runtime libs used by the tests / launcher. -conda install -c conda-forge -y cupy mpi4py nanobind blake3 sortedcontainers - -# PyTorch (pulls a matching cuda-toolkit + NCCL). -pip3 install torch - -# mscclpp build deps (used by `pip install .` of this repo). -pip install scikit-build-core nanobind setuptools_scm -``` - -Then build the EP extension (see [Build](#build)) — `pip install .` from -the repo root installs `mscclpp_ep_cpp.so` into the active env so the -test scripts can `import mscclpp.ep as ep`. - -Intranode (single node, 8 GPUs) — HT: - -```bash -MSCCLPP_EP_BENCH=1 \ -MSCCLPP_EP_BENCH_TOKENS=4096 MSCCLPP_EP_BENCH_HIDDEN=7168 \ -MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ -torchrun --nnodes=1 --nproc_per_node=8 \ +HWLOC_COMPONENTS=-gl \ +LD_LIBRARY_PATH=/usr/local/cuda/lib64 \ +torchrun --standalone --nproc_per_node=8 \ test/python/ep/test_intranode_multirank.py ``` -Intranode LL (single node, 8 GPUs), with no topology env: +The LL validation remains: ```bash +HWLOC_COMPONENTS=-gl \ +LD_LIBRARY_PATH=/usr/local/cuda/lib64 \ torchrun --standalone --nproc_per_node=8 \ - test/python/ep/test_low_latency_multirank.py \ - --num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256 \ - --cuda-graph + test/python/ep/test_low_latency_multirank.py ``` - -Internode HT (2 nodes × 8 GPUs), torchrun: - -```bash -# node 0 (master) -MSCCLPP_EP_BENCH=1 \ -MSCCLPP_EP_BENCH_TOKENS=4096 MSCCLPP_EP_BENCH_HIDDEN=7168 \ -MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ -torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 \ - --master_addr= --master_port=29600 \ - test/python/ep/test_internode_multirank.py - -# node 1 (worker) -MSCCLPP_EP_BENCH=1 \ -MSCCLPP_EP_BENCH_TOKENS=4096 MSCCLPP_EP_BENCH_HIDDEN=7168 \ -MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ -torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 \ - --master_addr= --master_port=29600 \ - test/python/ep/test_internode_multirank.py -``` - -If the bootstrap NIC is mis-detected (e.g. multi-homed hosts), pin -it explicitly: - -```bash -export NCCL_SOCKET_IFNAME= -export MSCCLPP_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME -export GLOO_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME -``` - -Internode HT via mpirun (NCCL-EP convention with NUMA binding): - -```bash -mpirun -np 16 --allow-run-as-root --hostfile \ - --bind-to numa \ - -x MSCCLPP_EP_BENCH=1 \ - -x MSCCLPP_EP_BENCH_TOKENS=4096 -x MSCCLPP_EP_BENCH_HIDDEN=7168 \ - -x MSCCLPP_EP_BENCH_EXPERTS=256 -x MSCCLPP_EP_BENCH_TOPK=8 \ - -x MASTER_ADDR= -x MASTER_PORT=29600 \ - bash -c 'export RANK=$OMPI_COMM_WORLD_RANK \ - WORLD_SIZE=$OMPI_COMM_WORLD_SIZE \ - LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK; \ - exec python3 test/python/ep/test_internode_multirank.py' -``` - -Cross-node LL is supported only when bootstrap reports one GPU IPC domain -(for example GB200 NVL72 with IMEX). The launch does not pass a domain size: - -```bash -mpirun -np 16 --allow-run-as-root --hostfile \ - --bind-to numa \ - -x MASTER_ADDR= -x MASTER_PORT=29600 \ - bash -c 'export RANK=$OMPI_COMM_WORLD_RANK \ - WORLD_SIZE=$OMPI_COMM_WORLD_SIZE \ - LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK; \ - exec python3 test/python/ep/test_low_latency_multirank.py \ - --num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256' -``` - -Add `-x NCCL_SOCKET_IFNAME= -x MSCCLPP_SOCKET_IFNAME= --x GLOO_SOCKET_IFNAME=` to the `mpirun` lines above only if the -default bootstrap NIC is wrong. These variables affect the launcher/bootstrap, -not LL topology detection. - -If Open MPI's own bootstrap misbehaves (e.g. UCX is mis-configured or -the host is multi-homed), force its control channel onto plain TCP -over the local management NIC by adding -`--mca pml ob1 --mca btl tcp,vader,self --mca btl_tcp_if_include ` -to `mpirun`. `` is the management NIC (e.g. `enP22p1s0f1` on -Azure GB200, `eno1`/`bond0` elsewhere) — it is **not** `eth0` on -GB200. - -### Benchmark mode - -Use the unified driver for NCCL-EP-style LL measurements: - -```bash -# Build the pure-C++ benchmark. -cmake -S test/python/ep -B test/python/ep/build \ - -DCMAKE_CUDA_ARCHITECTURES=90 -cmake --build test/python/ep/build -j 64 - -# BF16 rank-local. -python3 test/python/ep/run_ep_bench.py \ - --ep-lib mscclpp-cpp --nproc-per-node 8 \ - -t 128 -d 7168 -k 8 -e 256 -w 50 -i 100 \ - --dispatch-dtype bf16 --combine-mode rank_local_reduce - -# FP8 direct-send. -python3 test/python/ep/run_ep_bench.py \ - --ep-lib mscclpp-cpp --nproc-per-node 8 \ - -t 128 -d 7168 -k 8 -e 256 -w 50 -i 100 \ - --dispatch-dtype fp8_e4m3 --combine-mode direct_send -``` - -`ep_bench_ll.py` provides the high-level Python equivalent. The existing -multirank HT tests also expose their older env-controlled benchmark pass. - -HT multirank test benchmark env knobs: - -| Variable | Meaning | Default | -|-------------------------------|----------------------------------------|---------| -| `MSCCLPP_EP_BENCH` | Enable benchmark pass | `0` | -| `MSCCLPP_EP_BENCH_WARMUP` | Warmup iterations | `10` | -| `MSCCLPP_EP_BENCH_ITERS` | Timed iterations | `50` | -| `MSCCLPP_EP_BENCH_TOKENS` | Tokens per rank | test-specific | -| `MSCCLPP_EP_BENCH_HIDDEN` | Hidden dim | `7168` | -| `MSCCLPP_EP_BENCH_EXPERTS` | Total experts | test-specific | -| `MSCCLPP_EP_BENCH_TOPK` | top-k routing | `8` | - -## Implementation notes - -### Backend map - -`mscclpp_ep_cpp` exposes both backends in one Python extension: - -| Backend | Python API | C++ runtime | Kernel sources | Layout | -|---------|------------|-------------|----------------|--------| -| LL | `MoECommunicator(mode=MoEMode.LOW_LATENCY)` / `MoERuntime` | `moe_runtime.cc` | `low_latency/{dispatch,combine}.cu` | `EXPERT_MAJOR` | -| HT | `MoECommunicator(mode=MoEMode.HIGH_THROUGHPUT)` / `ExpertParallelRuntime` | `ht_runtime.cc` | `ht/kernels/*` | `FLAT` | - -Shared internal headers live in `include/`. - -The LL runtime uses one `low_latency_num_blocks` setting. Its default is 130: -dispatch launches 128 workers plus scheduler/notify blocks, while combine -launches 128 workers. `RANK_LOCAL_REDUCE` is the default combine mode; -`DIRECT_SEND` preserves bit-exact top-k reduction order. - -LL accepts BF16 input and can produce BF16 or FP8 E4M3 dispatch output. FP8 -uses one FP32 scale per 128 hidden elements. Expert output passed to combine is -always BF16. - -The runtime owns two fixed communication receive regions inside one symmetric -allocation: dispatch always writes `dispatchRecvBuffer`, while combine reads -that region and writes `combineRecvBuffer`. Buffer sizing and allocation remain -entirely inside C++; Python passes only the static workload dimensions. - -### LL dispatch payload - -Dispatch deduplicates routing by destination rank: a token routed to multiple -experts on the same rank sends its hidden data only once. The receiver expands -that payload into all matching local-expert rows. Each 128-byte-aligned payload -contains: - -```text -[data: BF16[hidden] or FP8_E4M3[hidden]] -[optional scales: FP32[hidden / 128]] -[topk expert ids: int32[num_topk]] -[topk weights: FP32[num_topk]] -[source token global index: int32] -``` - -Small rank/expert counts are sent as LL8 packets. After every payload TMA store -has completed, the sender publishes readiness through a lightweight -`BaseMemoryChannel` signal; the receiver waits on that signal before copying -data into expert-major output. - -The optimized kernels are instantiated for hidden sizes `4096`, `6656`, -`7168`, `8192`, and `9216`; other hidden sizes are rejected. FP8 E4M3 -currently fixes the scale block at 128. A future scale layout must use a -distinct `DispatchDataType`. - -### Current H100 performance - -Measured on 8×H100 with 128 tokens/rank, hidden 7168, top-k 8, 256 experts, -50 warmup iterations, and NCCL-EP-style random routing with masked entries. -No LL topology environment variables were set. - -CUDA Graph dispatch+combine E2E: - -| Dispatch format | Combine mode | E2E | -|---|---|---:| -| BF16 | rank-local reduce | 81.2 µs | -| BF16 | direct send | 100.2 µs | -| FP8 E4M3 | rank-local reduce | 73.1 µs | -| FP8 E4M3 | direct send | 93.5 µs | - -Pure-C++ paired benchmark (`dispatch → synchronize → combine → synchronize`): - -| Dispatch format | Combine mode | Dispatch | Combine | D+C | -|---|---|---:|---:|---:| -| BF16 | rank-local reduce | 45.38 µs | 45.66 µs | 99.33 µs | -| BF16 | direct send | 45.18 µs | 64.39 µs | 117.83 µs | -| FP8 E4M3 | rank-local reduce | 35.59 µs | 45.98 µs | 89.67 µs | -| FP8 E4M3 | direct send | 35.76 µs | 66.68 µs | 110.55 µs | - -With the same BF16 expert-major workload, NCCL-EP measured 57.99 µs dispatch, -75.19 µs combine, and 141.26 µs D+C. MSCCL++ rank-local measured 45.49 µs, -45.30 µs, and 98.74 µs respectively in that comparison run. - -Rank-local combine reduces expert rows on each destination rank before sending -one BF16 partial per source rank/token. It may differ by one BF16 ULP because -the reduction order changes. `DIRECT_SEND` preserves bit-exact top-k order. diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index b337b8ae..24403f86 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -10,22 +10,14 @@ // (uintptr_t) boundary so the module never links libtorch: // - MoERuntime: low-latency (LL) path. // - MoEHighThroughputRuntime (exposed as `ExpertParallelRuntime`): the -// high-throughput (HT) DeepEP-style path. Dynamic recv sizing uses an -// explicit two-phase API (intranode_notify_dispatch -> caller allocates -> -// intranode_dispatch); the caller passes device pointers (`tensor.data_ptr()`) +// intranode high-throughput (HT) path. Dynamic recv sizing uses an explicit +// two-phase API (notify_dispatch -> caller allocates -> dispatch); the caller +// passes device pointers (`tensor.data_ptr()`) // and sizes, exactly like the LL runtime. #include -#include -#include -#include #include -#include -#include -#include -#include -#include #include "api.cuh" #include "config.hpp" @@ -37,27 +29,10 @@ namespace nb = nanobind; namespace { -nb::bytes stringToBytes(const std::string& data) { return nb::bytes(data.data(), data.size()); } - void* ptr(uintptr_t address) { return reinterpret_cast(address); } cudaStream_t stream(uintptr_t address) { return reinterpret_cast(address); } -// Python bytes / bytearray -> std::string (raw bytes). -std::string bytesToString(nb::handle h) { - PyObject* o = h.ptr(); - if (PyBytes_Check(o)) { - char* buf = nullptr; - Py_ssize_t len = 0; - PyBytes_AsStringAndSize(o, &buf, &len); - return std::string(buf, static_cast(len)); - } - if (PyByteArray_Check(o)) { - return std::string(PyByteArray_AsString(o), static_cast(PyByteArray_Size(o))); - } - throw std::runtime_error("mscclpp_ep: expected bytes or bytearray"); -} - } // namespace NB_MODULE(mscclpp_ep_cpp, m) { @@ -121,145 +96,77 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_data_type"), nb::arg("mode"), nb::arg("num_blocks"), nb::arg("stream_ptr")); - // ========================================================================== - // High-throughput (HT) DeepEP-style backend: Config + MoEHighThroughputRuntime. - // Torch-free, raw-pointer boundary (uintptr_t == device pointer). Dynamic recv - // sizing uses the two-phase notify -> (caller allocates) -> dispatch API. - // ========================================================================== - nb::class_(m, "Config") - .def(nb::init(), nb::arg("num_sms") = 20, nb::arg("num_max_nvl_chunked_send_tokens") = 6, - nb::arg("num_max_nvl_chunked_recv_tokens") = 256, nb::arg("num_max_rdma_chunked_send_tokens") = 6, - nb::arg("num_max_rdma_chunked_recv_tokens") = 256) - .def("get_nvl_buffer_size_hint", &mscclpp::ep::Config::get_nvl_buffer_size_hint) - .def("get_rdma_buffer_size_hint", &mscclpp::ep::Config::get_rdma_buffer_size_hint); + .def(nb::init(), nb::arg("num_sms") = 20, nb::arg("num_max_nvl_chunked_send_tokens") = 6, + nb::arg("num_max_nvl_chunked_recv_tokens") = 256) + .def("get_nvl_buffer_size_hint", &mscclpp::ep::Config::get_nvl_buffer_size_hint); nb::class_(m, "ExpertParallelRuntime") - .def(nb::init(), nb::arg("rank"), nb::arg("num_ranks"), nb::arg("num_nvl_bytes"), - nb::arg("num_rdma_bytes")) + .def(nb::init(), nb::arg("comm"), + nb::arg("max_hidden_bytes"), nb::arg("config")) .def("is_available", &mscclpp::ep::MoEHighThroughputRuntime::isAvailable) .def("is_internode_available", &mscclpp::ep::MoEHighThroughputRuntime::isInternodeAvailable) - .def("get_num_rdma_ranks", &mscclpp::ep::MoEHighThroughputRuntime::getNumRdmaRanks) - .def("get_rdma_rank", &mscclpp::ep::MoEHighThroughputRuntime::getRdmaRank) - .def("get_root_rdma_rank", &mscclpp::ep::MoEHighThroughputRuntime::getRootRdmaRank) - .def("get_local_device_id", &mscclpp::ep::MoEHighThroughputRuntime::getLocalDeviceId) - .def("get_local_ipc_handle", - [](const mscclpp::ep::MoEHighThroughputRuntime& self) { return stringToBytes(self.getLocalIpcHandle()); }) - .def("create_unique_id", - [](const mscclpp::ep::MoEHighThroughputRuntime& self) { - auto uid = self.createUniqueId(); - return nb::bytes(reinterpret_cast(uid.data()), uid.size()); - }) - .def("connect", - [](mscclpp::ep::MoEHighThroughputRuntime& self, nb::object data) { - std::string s = bytesToString(data); - mscclpp::UniqueId uid; - if (s.size() != uid.size()) { - throw std::runtime_error("mscclpp_ep_cpp.ExpertParallelRuntime.connect: UniqueId size mismatch"); - } - std::memcpy(uid.data(), s.data(), s.size()); - self.connect(uid); - }) .def( - "sync", - [](mscclpp::ep::MoEHighThroughputRuntime& self, const std::vector& device_ids, - nb::sequence all_gathered_handles, nb::object root_unique_id_opt) { - std::vector> handles; - for (nb::handle h : all_gathered_handles) { - if (h.is_none()) - handles.emplace_back(std::nullopt); - else - handles.emplace_back(bytesToString(h)); - } - std::optional root_uid = root_unique_id_opt.is_none() - ? std::nullopt - : std::optional(bytesToString(root_unique_id_opt)); - self.sync(device_ids, handles, root_uid); - }, - nb::arg("device_ids"), nb::arg("all_gathered_handles").none(), nb::arg("root_unique_id").none()) - .def( - "get_dispatch_layout", + "layout", [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t num_tokens_per_rank_ptr, - uintptr_t num_tokens_per_rdma_rank_ptr, uintptr_t num_tokens_per_expert_ptr, - uintptr_t is_token_in_rank_ptr, uintptr_t topk_idx_ptr, int num_tokens, int num_topk, int num_experts, - uintptr_t stream_ptr) { - self.getDispatchLayout(reinterpret_cast(ptr(num_tokens_per_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_rdma_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_expert_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), - reinterpret_cast(ptr(topk_idx_ptr)), num_tokens, num_topk, - num_experts, stream(stream_ptr)); + uintptr_t num_tokens_per_expert_ptr, uintptr_t is_token_in_rank_ptr, uintptr_t topk_idx_ptr, + int num_tokens, int num_topk, int num_experts, uintptr_t stream_ptr) { + self.layout(reinterpret_cast(ptr(num_tokens_per_rank_ptr)), + reinterpret_cast(ptr(num_tokens_per_expert_ptr)), + reinterpret_cast(ptr(is_token_in_rank_ptr)), + reinterpret_cast(ptr(topk_idx_ptr)), num_tokens, num_topk, num_experts, + stream(stream_ptr)); }, - nb::arg("num_tokens_per_rank_ptr"), nb::arg("num_tokens_per_rdma_rank_ptr"), - nb::arg("num_tokens_per_expert_ptr"), nb::arg("is_token_in_rank_ptr"), nb::arg("topk_idx_ptr"), - nb::arg("num_tokens"), nb::arg("num_topk"), nb::arg("num_experts"), nb::arg("stream_ptr")) - .def( - "get_intranode_dispatch_num_channels", - [](const mscclpp::ep::MoEHighThroughputRuntime& self, int x_element_size, const mscclpp::ep::Config& config) { - return self.getIntranodeDispatchNumChannels(x_element_size, config); - }) - .def("resolve_intranode_recv_x_buffer", - [](const mscclpp::ep::MoEHighThroughputRuntime& self, int num_recv_tokens, int hidden, int x_element_size, - const mscclpp::ep::Config& config) -> uintptr_t { + nb::arg("num_tokens_per_rank_ptr"), nb::arg("num_tokens_per_expert_ptr"), nb::arg("is_token_in_rank_ptr"), + nb::arg("topk_idx_ptr"), nb::arg("num_tokens"), nb::arg("num_topk"), nb::arg("num_experts"), + nb::arg("stream_ptr")) + .def("get_dispatch_num_channels", [](const mscclpp::ep::MoEHighThroughputRuntime& self, + int x_element_size) { return self.getDispatchNumChannels(x_element_size); }) + .def("resolve_recv_x_buffer", + [](const mscclpp::ep::MoEHighThroughputRuntime& self, int num_tokens, int num_recv_tokens, int hidden, + int x_element_size) -> uintptr_t { return reinterpret_cast( - self.resolveIntranodeRecvXBuffer(num_recv_tokens, hidden, x_element_size, config)); + self.resolveRecvXBuffer(num_tokens, num_recv_tokens, hidden, x_element_size)); }) - .def("resolve_internode_recv_x_buffer", - [](const mscclpp::ep::MoEHighThroughputRuntime& self, int num_recv_tokens, int hidden, int x_element_size, - const mscclpp::ep::Config& config) -> uintptr_t { - return reinterpret_cast( - self.resolveInternodeRecvXBuffer(num_recv_tokens, hidden, x_element_size, config)); - }) - .def("get_internode_dispatch_num_channels", - [](const mscclpp::ep::MoEHighThroughputRuntime& self, const mscclpp::ep::Config& config) { - return self.getInternodeDispatchNumChannels(config); - }) - .def("get_source_meta_bytes", - [](const mscclpp::ep::MoEHighThroughputRuntime& self) { return self.getSourceMetaBytes(); }) - .def("get_num_max_nvl_peers", - [](const mscclpp::ep::MoEHighThroughputRuntime& self) { return self.getNumMaxNvlPeers(); }) .def( - "intranode_notify_dispatch", + "notify_dispatch", [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t rank_prefix_matrix_ptr, uintptr_t channel_prefix_matrix_ptr, uintptr_t num_recv_tokens_per_expert_ptr, uintptr_t num_tokens_per_rank_ptr, uintptr_t num_tokens_per_expert_ptr, uintptr_t is_token_in_rank_ptr, - int num_tokens, int num_experts, int x_element_size, int expert_alignment, - const mscclpp::ep::Config& config, uintptr_t stream_ptr) { - return self.intranodeNotifyDispatch(reinterpret_cast(ptr(rank_prefix_matrix_ptr)), - reinterpret_cast(ptr(channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(num_recv_tokens_per_expert_ptr)), - reinterpret_cast(ptr(num_tokens_per_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_expert_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), num_tokens, - num_experts, x_element_size, expert_alignment, config, - stream(stream_ptr)); + int num_tokens, int num_experts, int x_element_size, int expert_alignment, uintptr_t stream_ptr) { + return self.notifyDispatch(reinterpret_cast(ptr(rank_prefix_matrix_ptr)), + reinterpret_cast(ptr(channel_prefix_matrix_ptr)), + reinterpret_cast(ptr(num_recv_tokens_per_expert_ptr)), + reinterpret_cast(ptr(num_tokens_per_rank_ptr)), + reinterpret_cast(ptr(num_tokens_per_expert_ptr)), + reinterpret_cast(ptr(is_token_in_rank_ptr)), num_tokens, + num_experts, x_element_size, expert_alignment, stream(stream_ptr)); }, nb::arg("rank_prefix_matrix_ptr"), nb::arg("channel_prefix_matrix_ptr"), nb::arg("num_recv_tokens_per_expert_ptr"), nb::arg("num_tokens_per_rank_ptr"), nb::arg("num_tokens_per_expert_ptr"), nb::arg("is_token_in_rank_ptr"), nb::arg("num_tokens"), - nb::arg("num_experts"), nb::arg("x_element_size"), nb::arg("expert_alignment"), nb::arg("config"), - nb::arg("stream_ptr")) + nb::arg("num_experts"), nb::arg("x_element_size"), nb::arg("expert_alignment"), nb::arg("stream_ptr")) .def( - "intranode_dispatch", + "dispatch", [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t recv_x_ptr, uintptr_t recv_x_scales_ptr, uintptr_t recv_topk_idx_ptr, uintptr_t recv_topk_weights_ptr, uintptr_t recv_src_idx_ptr, uintptr_t send_head_ptr, uintptr_t recv_channel_prefix_matrix_ptr, uintptr_t x_ptr, uintptr_t x_scales_ptr, uintptr_t topk_idx_ptr, uintptr_t topk_weights_ptr, uintptr_t is_token_in_rank_ptr, uintptr_t rank_prefix_matrix_ptr, uintptr_t channel_prefix_matrix_ptr, int num_tokens, int hidden, int num_topk, int num_scales, int num_experts, int x_element_size, int num_recv_tokens, bool cached_mode, - const mscclpp::ep::Config& config, uintptr_t stream_ptr) { - self.intranodeDispatch( - ptr(recv_x_ptr), reinterpret_cast(ptr(recv_x_scales_ptr)), - reinterpret_cast(ptr(recv_topk_idx_ptr)), - reinterpret_cast(ptr(recv_topk_weights_ptr)), reinterpret_cast(ptr(recv_src_idx_ptr)), - reinterpret_cast(ptr(send_head_ptr)), reinterpret_cast(ptr(recv_channel_prefix_matrix_ptr)), - ptr(x_ptr), reinterpret_cast(ptr(x_scales_ptr)), - reinterpret_cast(ptr(topk_idx_ptr)), - reinterpret_cast(ptr(topk_weights_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), - reinterpret_cast(ptr(rank_prefix_matrix_ptr)), - reinterpret_cast(ptr(channel_prefix_matrix_ptr)), num_tokens, hidden, num_topk, num_scales, - num_experts, x_element_size, num_recv_tokens, cached_mode, config, stream(stream_ptr)); + uintptr_t stream_ptr) { + self.dispatch(ptr(recv_x_ptr), reinterpret_cast(ptr(recv_x_scales_ptr)), + reinterpret_cast(ptr(recv_topk_idx_ptr)), + reinterpret_cast(ptr(recv_topk_weights_ptr)), + reinterpret_cast(ptr(recv_src_idx_ptr)), reinterpret_cast(ptr(send_head_ptr)), + reinterpret_cast(ptr(recv_channel_prefix_matrix_ptr)), ptr(x_ptr), + reinterpret_cast(ptr(x_scales_ptr)), + reinterpret_cast(ptr(topk_idx_ptr)), + reinterpret_cast(ptr(topk_weights_ptr)), + reinterpret_cast(ptr(is_token_in_rank_ptr)), + reinterpret_cast(ptr(rank_prefix_matrix_ptr)), + reinterpret_cast(ptr(channel_prefix_matrix_ptr)), num_tokens, hidden, num_topk, + num_scales, num_experts, x_element_size, num_recv_tokens, cached_mode, stream(stream_ptr)); }, nb::arg("recv_x_ptr"), nb::arg("recv_x_scales_ptr"), nb::arg("recv_topk_idx_ptr"), nb::arg("recv_topk_weights_ptr"), nb::arg("recv_src_idx_ptr"), nb::arg("send_head_ptr"), @@ -267,117 +174,24 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::arg("topk_weights_ptr"), nb::arg("is_token_in_rank_ptr"), nb::arg("rank_prefix_matrix_ptr"), nb::arg("channel_prefix_matrix_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("num_scales"), nb::arg("num_experts"), nb::arg("x_element_size"), nb::arg("num_recv_tokens"), - nb::arg("cached_mode"), nb::arg("config"), nb::arg("stream_ptr")) + nb::arg("cached_mode"), nb::arg("stream_ptr")) .def( - "intranode_combine", + "combine", [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t combined_x_ptr, uintptr_t combined_topk_weights_ptr, uintptr_t x_ptr, uintptr_t topk_weights_ptr, uintptr_t src_idx_ptr, uintptr_t rank_prefix_matrix_ptr, uintptr_t channel_prefix_matrix_ptr, uintptr_t send_head_ptr, int num_tokens, int num_recv_tokens, - int hidden, int num_topk, int x_element_size, int ring_num_channels, const mscclpp::ep::Config& config, - uintptr_t stream_ptr) { - self.intranodeCombine(ptr(combined_x_ptr), reinterpret_cast(ptr(combined_topk_weights_ptr)), - ptr(x_ptr), reinterpret_cast(ptr(topk_weights_ptr)), - reinterpret_cast(ptr(src_idx_ptr)), - reinterpret_cast(ptr(rank_prefix_matrix_ptr)), - reinterpret_cast(ptr(channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(send_head_ptr)), num_tokens, num_recv_tokens, hidden, - num_topk, x_element_size, ring_num_channels, config, stream(stream_ptr)); + int hidden, int num_topk, int x_element_size, int ring_num_channels, uintptr_t stream_ptr) { + self.combine(ptr(combined_x_ptr), reinterpret_cast(ptr(combined_topk_weights_ptr)), ptr(x_ptr), + reinterpret_cast(ptr(topk_weights_ptr)), + reinterpret_cast(ptr(src_idx_ptr)), + reinterpret_cast(ptr(rank_prefix_matrix_ptr)), + reinterpret_cast(ptr(channel_prefix_matrix_ptr)), + reinterpret_cast(ptr(send_head_ptr)), num_tokens, num_recv_tokens, hidden, + num_topk, x_element_size, ring_num_channels, stream(stream_ptr)); }, nb::arg("combined_x_ptr"), nb::arg("combined_topk_weights_ptr"), nb::arg("x_ptr"), nb::arg("topk_weights_ptr"), nb::arg("src_idx_ptr"), nb::arg("rank_prefix_matrix_ptr"), nb::arg("channel_prefix_matrix_ptr"), nb::arg("send_head_ptr"), nb::arg("num_tokens"), nb::arg("num_recv_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("x_element_size"), - nb::arg("ring_num_channels"), nb::arg("config"), nb::arg("stream_ptr")) - .def( - "internode_notify_dispatch", - [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t rdma_channel_prefix_matrix_ptr, - uintptr_t recv_rdma_rank_prefix_sum_ptr, uintptr_t gbl_channel_prefix_matrix_ptr, - uintptr_t recv_gbl_rank_prefix_sum_ptr, uintptr_t num_recv_tokens_per_expert_ptr, - uintptr_t num_rdma_recv_tokens_ptr, uintptr_t num_tokens_per_rank_ptr, - uintptr_t num_tokens_per_rdma_rank_ptr, uintptr_t num_tokens_per_expert_ptr, - uintptr_t is_token_in_rank_ptr, int num_tokens, int num_experts, int hidden, int num_scales, int num_topk, - int x_element_size, int expert_alignment, const mscclpp::ep::Config& config, uintptr_t stream_ptr) { - return self.internodeNotifyDispatch(reinterpret_cast(ptr(rdma_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(recv_rdma_rank_prefix_sum_ptr)), - reinterpret_cast(ptr(gbl_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(recv_gbl_rank_prefix_sum_ptr)), - reinterpret_cast(ptr(num_recv_tokens_per_expert_ptr)), - reinterpret_cast(ptr(num_rdma_recv_tokens_ptr)), - reinterpret_cast(ptr(num_tokens_per_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_rdma_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_expert_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), num_tokens, - num_experts, hidden, num_scales, num_topk, x_element_size, - expert_alignment, config, stream(stream_ptr)); - }, - nb::arg("rdma_channel_prefix_matrix_ptr"), nb::arg("recv_rdma_rank_prefix_sum_ptr"), - nb::arg("gbl_channel_prefix_matrix_ptr"), nb::arg("recv_gbl_rank_prefix_sum_ptr"), - nb::arg("num_recv_tokens_per_expert_ptr"), nb::arg("num_rdma_recv_tokens_ptr"), - nb::arg("num_tokens_per_rank_ptr"), nb::arg("num_tokens_per_rdma_rank_ptr"), - nb::arg("num_tokens_per_expert_ptr"), nb::arg("is_token_in_rank_ptr"), nb::arg("num_tokens"), - nb::arg("num_experts"), nb::arg("hidden"), nb::arg("num_scales"), nb::arg("num_topk"), - nb::arg("x_element_size"), nb::arg("expert_alignment"), nb::arg("config"), nb::arg("stream_ptr")) - .def( - "internode_dispatch", - [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t recv_x_ptr, uintptr_t recv_x_scales_ptr, - uintptr_t recv_topk_idx_ptr, uintptr_t recv_topk_weights_ptr, uintptr_t recv_src_meta_ptr, - uintptr_t recv_rdma_channel_prefix_matrix_ptr, uintptr_t recv_gbl_channel_prefix_matrix_ptr, - uintptr_t send_rdma_head_ptr, uintptr_t send_nvl_head_ptr, uintptr_t x_ptr, uintptr_t x_scales_ptr, - uintptr_t topk_idx_ptr, uintptr_t topk_weights_ptr, uintptr_t is_token_in_rank_ptr, - uintptr_t rdma_channel_prefix_matrix_ptr, uintptr_t recv_rdma_rank_prefix_sum_ptr, - uintptr_t gbl_channel_prefix_matrix_ptr, uintptr_t recv_gbl_rank_prefix_sum_ptr, int num_tokens, - int hidden, int num_topk, int num_scales, int num_experts, int x_element_size, int num_recv_tokens, - int num_rdma_recv_tokens, bool cached_mode, const mscclpp::ep::Config& config, uintptr_t stream_ptr) { - self.internodeDispatch(ptr(recv_x_ptr), reinterpret_cast(ptr(recv_x_scales_ptr)), - reinterpret_cast(ptr(recv_topk_idx_ptr)), - reinterpret_cast(ptr(recv_topk_weights_ptr)), ptr(recv_src_meta_ptr), - reinterpret_cast(ptr(recv_rdma_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(recv_gbl_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(send_rdma_head_ptr)), - reinterpret_cast(ptr(send_nvl_head_ptr)), ptr(x_ptr), - reinterpret_cast(ptr(x_scales_ptr)), - reinterpret_cast(ptr(topk_idx_ptr)), - reinterpret_cast(ptr(topk_weights_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), - reinterpret_cast(ptr(rdma_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(recv_rdma_rank_prefix_sum_ptr)), - reinterpret_cast(ptr(gbl_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(recv_gbl_rank_prefix_sum_ptr)), num_tokens, hidden, - num_topk, num_scales, num_experts, x_element_size, num_recv_tokens, - num_rdma_recv_tokens, cached_mode, config, stream(stream_ptr)); - }, - nb::arg("recv_x_ptr"), nb::arg("recv_x_scales_ptr"), nb::arg("recv_topk_idx_ptr"), - nb::arg("recv_topk_weights_ptr"), nb::arg("recv_src_meta_ptr"), - nb::arg("recv_rdma_channel_prefix_matrix_ptr"), nb::arg("recv_gbl_channel_prefix_matrix_ptr"), - nb::arg("send_rdma_head_ptr"), nb::arg("send_nvl_head_ptr"), nb::arg("x_ptr"), nb::arg("x_scales_ptr"), - nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("is_token_in_rank_ptr"), - nb::arg("rdma_channel_prefix_matrix_ptr"), nb::arg("recv_rdma_rank_prefix_sum_ptr"), - nb::arg("gbl_channel_prefix_matrix_ptr"), nb::arg("recv_gbl_rank_prefix_sum_ptr"), nb::arg("num_tokens"), - nb::arg("hidden"), nb::arg("num_topk"), nb::arg("num_scales"), nb::arg("num_experts"), - nb::arg("x_element_size"), nb::arg("num_recv_tokens"), nb::arg("num_rdma_recv_tokens"), - nb::arg("cached_mode"), nb::arg("config"), nb::arg("stream_ptr")) - .def( - "internode_combine", - [](mscclpp::ep::MoEHighThroughputRuntime& self, uintptr_t combined_x_ptr, uintptr_t combined_topk_weights_ptr, - uintptr_t x_ptr, uintptr_t topk_weights_ptr, uintptr_t src_meta_ptr, - uintptr_t is_combined_token_in_rank_ptr, uintptr_t rdma_channel_prefix_matrix_ptr, - uintptr_t rdma_rank_prefix_sum_ptr, uintptr_t gbl_channel_prefix_matrix_ptr, - uintptr_t combined_rdma_head_ptr, uintptr_t combined_nvl_head_ptr, int num_tokens, int num_combined_tokens, - int hidden, int num_topk, int x_element_size, const mscclpp::ep::Config& config, uintptr_t stream_ptr) { - self.internodeCombine(ptr(combined_x_ptr), reinterpret_cast(ptr(combined_topk_weights_ptr)), - ptr(x_ptr), reinterpret_cast(ptr(topk_weights_ptr)), ptr(src_meta_ptr), - reinterpret_cast(ptr(is_combined_token_in_rank_ptr)), - reinterpret_cast(ptr(rdma_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(rdma_rank_prefix_sum_ptr)), - reinterpret_cast(ptr(gbl_channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(combined_rdma_head_ptr)), - reinterpret_cast(ptr(combined_nvl_head_ptr)), num_tokens, - num_combined_tokens, hidden, num_topk, x_element_size, config, stream(stream_ptr)); - }, - nb::arg("combined_x_ptr"), nb::arg("combined_topk_weights_ptr"), nb::arg("x_ptr"), - nb::arg("topk_weights_ptr"), nb::arg("src_meta_ptr"), nb::arg("is_combined_token_in_rank_ptr"), - nb::arg("rdma_channel_prefix_matrix_ptr"), nb::arg("rdma_rank_prefix_sum_ptr"), - nb::arg("gbl_channel_prefix_matrix_ptr"), nb::arg("combined_rdma_head_ptr"), nb::arg("combined_nvl_head_ptr"), - nb::arg("num_tokens"), nb::arg("num_combined_tokens"), nb::arg("hidden"), nb::arg("num_topk"), - nb::arg("x_element_size"), nb::arg("config"), nb::arg("stream_ptr")); + nb::arg("ring_num_channels"), nb::arg("stream_ptr")); } diff --git a/src/ext/ep/ht/kernels/buffer.cuh b/src/ext/ep/ht/buffer.cuh similarity index 98% rename from src/ext/ep/ht/kernels/buffer.cuh rename to src/ext/ep/ht/buffer.cuh index 0ac87a64..8110ec72 100644 --- a/src/ext/ep/ht/kernels/buffer.cuh +++ b/src/ext/ep/ht/buffer.cuh @@ -2,6 +2,8 @@ // Licensed under the MIT License. #pragma once +// Device-side buffer views for HT ring storage. + #include #include "constants.cuh" diff --git a/src/ext/ep/ht/config.hpp b/src/ext/ep/ht/config.hpp index 71bc1289..59391bb8 100644 --- a/src/ext/ep/ht/config.hpp +++ b/src/ext/ep/ht/config.hpp @@ -2,143 +2,81 @@ // Licensed under the MIT License. #pragma once -#include +#include +#include +#include #include "../config.hpp" -#include "api.cuh" +#include "constants.cuh" #include "exception.cuh" namespace mscclpp { namespace ep { -// inc7: resolve the flat all-sender dispatch channel count (== block count). -// Default num_sms/2; MSCCLPP_EP_DISPATCH_NSM (flat path only) overrides it, clamped -// to [1, num_sms]. The flat path launches one sender block per channel (no forwarder), -// so it can use the FULL SM budget with a SINGLE knob. Non-flat/unset -> num_sms/2. -inline int ep_flat_dispatch_channels(int num_sms) { - const char* f = std::getenv("MSCCLPP_EP_FLAT"); - const char* d = std::getenv("MSCCLPP_EP_DIRECT"); - if (!(f != nullptr && std::atoi(f) != 0 && d != nullptr && std::atoi(d) != 0)) return num_sms / 2; - const char* e = std::getenv("MSCCLPP_EP_DISPATCH_NSM"); - if (e == nullptr) return num_sms / 2; - const int n = std::atoi(e); - if (n < 1) return num_sms / 2; - return n < num_sms ? n : num_sms; -} - -// Channel count the RDMA/NVL buffers must hold: max of the flat dispatch channels -// and the default num_sms/2 (combine grid). Keeps per-channel offsets in bounds when -// DISPATCH_NSM raises/lowers the dispatch channel count relative to num_sms/2. -inline int ep_buffer_channels(int num_sms) { - const int dc = ep_flat_dispatch_channels(num_sms); - const int half = num_sms / 2; - return dc > half ? dc : half; -} - struct Config { + static constexpr int MaxTopk = 32; + static constexpr int MaxScales = 128; + static constexpr int RecvPoolMaxTokens = 65536; + static constexpr int64_t RecvPoolMaxHiddenBytes = 16384; + static constexpr int64_t RecvPoolMetaBytes = 128; + int num_sms; int num_max_nvl_chunked_send_tokens; int num_max_nvl_chunked_recv_tokens; - int num_max_rdma_chunked_send_tokens; - int num_max_rdma_chunked_recv_tokens; - Config(int num_sms, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, - int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens) + Config(int num_sms, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens) : num_sms(num_sms), num_max_nvl_chunked_send_tokens(num_max_nvl_chunked_send_tokens), - num_max_nvl_chunked_recv_tokens(num_max_nvl_chunked_recv_tokens), - num_max_rdma_chunked_send_tokens(num_max_rdma_chunked_send_tokens), - num_max_rdma_chunked_recv_tokens(num_max_rdma_chunked_recv_tokens) { - EP_HOST_ASSERT(num_sms >= 0); + num_max_nvl_chunked_recv_tokens(num_max_nvl_chunked_recv_tokens) { + EP_HOST_ASSERT(num_sms > 0 and num_sms % 2 == 0); EP_HOST_ASSERT(num_max_nvl_chunked_send_tokens > 0 and num_max_nvl_chunked_recv_tokens > 0); EP_HOST_ASSERT(num_max_nvl_chunked_send_tokens < num_max_nvl_chunked_recv_tokens); - EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens > 0 and num_max_rdma_chunked_recv_tokens > 0); - - // Ceil up RDMA buffer size - this->num_max_rdma_chunked_recv_tokens = - configAlign(num_max_rdma_chunked_recv_tokens, num_max_rdma_chunked_send_tokens); - EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens < num_max_rdma_chunked_recv_tokens); - // NOTES: this assertion is related to RDMA lazy head update, we must ensure senders always have space to push - EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens <= num_max_rdma_chunked_recv_tokens / 2); } - size_t get_nvl_base_bytes(size_t hidden_bytes, int num_ranks) const { - // Below are some assumptions - // TODO: add assertions - constexpr int kNumMaxTopK = 128; - constexpr int kNumMaxScales = 128; - EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0); - EP_HOST_ASSERT(num_ranks <= NUM_MAX_NVL_PEERS or num_sms % 2 == 0); - const auto num_rdma_ranks = std::max(num_ranks / NUM_MAX_NVL_PEERS, 1); - const auto num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS); - const int num_channels = ep_buffer_channels(num_sms); - - size_t num_bytes = 0; - num_bytes += num_channels * num_nvl_ranks * (2 * num_rdma_ranks + 3) * sizeof(int); - num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * hidden_bytes; - num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * internode::get_source_meta_bytes(); - num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(int64_t); - num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxTopK * sizeof(float); - num_bytes += num_channels * num_nvl_ranks * num_max_nvl_chunked_recv_tokens * kNumMaxScales * sizeof(float); - num_bytes = ((num_bytes + 127) / 128) * 128; - return num_bytes; - } - -#ifdef EP_DISPATCH_NCCLEP - static constexpr int kEpRecvPoolMaxTokens = 65536; - static constexpr int64_t kEpRecvPoolMaxHiddenBytes = 16384; // worst-case per-token bytes (hidden<=8192 bf16) - // inc6 (flat all-sender dispatch): per-token meta region appended AFTER the - // worst-case hidden region. The flat sender writes SourceMeta + scales + topk - // straight into the dest pool meta region; a local consumer copies it into the - // recv_* tensors. 128B/token (allocated unconditionally; only touched under flat). - static constexpr int64_t kEpRecvPoolMetaBytes = 128; // per-token meta slot (128B aligned) - size_t get_recv_pool_header_bytes(int num_ranks) const { - return ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - } - // Byte offset (from pool base) where the inc6 meta region starts: after the - // header and the full worst-case hidden region, so it never overlaps runtime - // hidden tokens regardless of the runtime hidden dim. - static size_t get_recv_pool_meta_base(int num_ranks) { - size_t header = ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - size_t hidden = static_cast(kEpRecvPoolMaxTokens) * static_cast(kEpRecvPoolMaxHiddenBytes); - return ((header + hidden + 127) / 128) * 128; - } - static size_t recv_pool_bytes_static(int num_ranks) { - size_t meta_base = get_recv_pool_meta_base(num_ranks); - size_t b = meta_base + static_cast(kEpRecvPoolMaxTokens) * static_cast(kEpRecvPoolMetaBytes); - return ((b + 127) / 128) * 128; - } -#endif - size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) const { - return get_nvl_base_bytes(hidden_bytes, num_ranks); + EP_HOST_ASSERT(hidden_bytes > 0); + EP_HOST_ASSERT(num_ranks == 2 or num_ranks == 4 or num_ranks == 8); + + const size_t ranks = static_cast(num_ranks); + const size_t recv_tokens = static_cast(num_max_nvl_chunked_recv_tokens); + const size_t ring_channels = static_cast(num_sms / 2); + const size_t all_sender_channels = static_cast(num_sms); + const size_t prefix_bytes = ranks * ranks * sizeof(int); + const size_t expert_scratch_bytes = ranks * NUM_MAX_LOCAL_EXPERTS * sizeof(int); + + const size_t dispatch_metadata_bytes = 4 * ring_channels * ranks * sizeof(int); + const size_t dispatch_token_bytes = + hidden_bytes + sizeof(int) + MaxTopk * sizeof(int64_t) + MaxTopk * sizeof(float) + MaxScales * sizeof(float); + const size_t dispatch_bytes = + prefix_bytes + std::max(expert_scratch_bytes, + dispatch_metadata_bytes + ring_channels * ranks * recv_tokens * dispatch_token_bytes); + + const size_t combine_metadata_bytes = 2 * ring_channels * ranks * sizeof(int); + const size_t combine_token_bytes = hidden_bytes + sizeof(int) + MaxTopk * sizeof(float); + const size_t combine_bytes = combine_metadata_bytes + ring_channels * ranks * recv_tokens * combine_token_bytes; + + const size_t all_sender_bytes = + prefix_bytes + std::max(expert_scratch_bytes, 4 * all_sender_channels * ranks * sizeof(int)); + return configAlign(std::max({dispatch_bytes, combine_bytes, all_sender_bytes}), NUM_BUFFER_ALIGNMENT_BYTES); } - size_t get_rdma_buffer_size_hint(int64_t hidden_bytes, int num_ranks) const { - // Legacy mode - if (num_ranks <= NUM_MAX_NVL_PEERS) return 0; + static size_t get_recv_pool_header_bytes(int num_ranks) { + return configAlign(static_cast(num_ranks) * sizeof(int), NUM_BUFFER_ALIGNMENT_BYTES); + } - // Below are some assumptions - // TODO: add assertions - constexpr int kNumMaxTopK = 128; - constexpr int kNumMaxScales = 128; - EP_HOST_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0); - EP_HOST_ASSERT(num_sms % 2 == 0); - const int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; - const int num_channels = ep_buffer_channels(num_sms); + static size_t get_recv_pool_meta_base(int num_ranks) { + const size_t hidden_bytes = static_cast(RecvPoolMaxTokens) * static_cast(RecvPoolMaxHiddenBytes); + return configAlign(get_recv_pool_header_bytes(num_ranks) + hidden_bytes, NUM_BUFFER_ALIGNMENT_BYTES); + } - size_t num_bytes = 0; - num_bytes += num_channels * num_rdma_ranks * (NUM_MAX_NVL_PEERS * 2 + 2) * 2 * sizeof(int); - num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * hidden_bytes * 2; - num_bytes += - num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * internode::get_source_meta_bytes() * 2; - num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(int64_t) * 2; - num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxTopK * sizeof(float) * 2; - num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * kNumMaxScales * sizeof(float) * 2; - num_bytes += num_channels * num_rdma_ranks * num_max_rdma_chunked_recv_tokens * sizeof(int4) * 2; - num_bytes += num_channels * num_rdma_ranks * sizeof(uint64_t) * 2; - num_bytes = ((num_bytes + 127) / 128) * 128; - return num_bytes; + static size_t get_recv_pool_hidden_bytes(int num_ranks) { + return get_recv_pool_meta_base(num_ranks) - get_recv_pool_header_bytes(num_ranks); + } + + static size_t recv_pool_bytes_static(int num_ranks) { + const size_t bytes = + get_recv_pool_meta_base(num_ranks) + static_cast(RecvPoolMaxTokens) * RecvPoolMetaBytes; + return configAlign(bytes, NUM_BUFFER_ALIGNMENT_BYTES); } }; diff --git a/src/ext/ep/ht/kernels/intranode_kernel.cu b/src/ext/ep/ht/intranode_kernel.cu similarity index 98% rename from src/ext/ep/ht/kernels/intranode_kernel.cu rename to src/ext/ep/ht/intranode_kernel.cu index 19c93e5a..ffcfdb33 100644 --- a/src/ext/ep/ht/kernels/intranode_kernel.cu +++ b/src/ext/ep/ht/intranode_kernel.cu @@ -1,5 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +// +// HT dispatch and combine kernels for directly mapped peers. + #include #include "buffer.cuh" @@ -510,7 +513,7 @@ void dispatch(void* recv_x, float* recv_x_scales, int* recv_src_idx, int64_t* re // combine gather map, and pack per-token metadata (src_idx, rebased topk_idx, // topk_weights, scales) into the destination pool's META region at the final recv // slot. A separate intranode_meta_drain kernel then unpacks the local pool META -// region into the recv_* output tensors (mirrors the internode flat_meta_drain). +// region into the recv_* output tensors. // No ring, no head/tail flow control, no receiver. Pairs with the TMA combine, // which is token-parallel and ignores the channel prefix matrix. template @@ -592,7 +595,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) } // Unpack the local recv-pool META region (filled by the all-sender dispatch) into the -// recv_* output tensors. One thread per recv token. Mirrors internode flat_meta_drain. +// recv_* output tensors. One thread per recv token. __global__ void intranode_meta_drain_kernel(const uint8_t* __restrict__ pool_base, int64_t meta_base, int num_recv_tokens, int* __restrict__ recv_src_idx, int64_t* __restrict__ recv_topk_idx, float* __restrict__ recv_topk_weights, @@ -640,8 +643,9 @@ void dispatch_allsender(int* send_head, const void* x, const int64_t* topk_idx, constexpr int kNumThreads = 512; EP_HOST_ASSERT(recv_pool_ptrs != nullptr); // Meta slot must hold src_idx + topk_idx(int) + topk_weights(float) + scales(float). - EP_HOST_ASSERT(static_cast(sizeof(int)) + static_cast(num_topk) * (sizeof(int) + sizeof(float)) + - static_cast(num_scales) * sizeof(float) <= + EP_HOST_ASSERT(static_cast(sizeof(int)) + + static_cast(num_topk) * static_cast(sizeof(int) + sizeof(float)) + + static_cast(num_scales) * static_cast(sizeof(float)) <= meta_slot_bytes); #define DISPATCH_ALLSENDER_LAUNCH_CASE(ranks) \ LAUNCH_KERNEL(&cfg, (dispatch_allsender), send_head, reinterpret_cast(x), topk_idx, \ @@ -985,11 +989,10 @@ __global__ void __launch_bounds__(kNumThreads, 1) } // --------------------------------------------------------------------------- -// Intranode TMA-staged direct-gather combine (port of internode -// combine_flat_gather_tma to the single-node NVLink/IPC path). +// Intranode TMA-staged direct-gather combine for the single-node peer-mapped path. // -// Mirrors the internode flat direct-gather: for each combined output token, it -// discovers the contributing ranks and gathers each contributor's hidden row +// For each combined output token, it discovers the contributing ranks and +// gathers each contributor's hidden row // straight from that rank's IPC-mapped recv-output pool (recv_pool_ptrs[r] at // slot ep_combine_recv_idx[t,r]) through a kStages-deep cp.async.bulk (TMA) SMEM // pipeline, reduces from SMEM, and writes the summed row to combined_x. No @@ -999,7 +1002,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) // Contributor discovery uses send_head (>=0 == token routed to that rank during // dispatch); the per-(token,rank) recv-pool slot comes from ep_combine_recv_idx, // which the dispatch sender-direct path fills. combined_topk_weights is zeroed -// (the intranode test, like the internode flat path, validates only combined_x). +// (the current intranode test validates only combined_x). #ifndef EP_ICMB_TMA_CHUNK_INT4 #define EP_ICMB_TMA_CHUNK_INT4 64 // hidden chunk in int4 (1KB TMA descriptors) #endif diff --git a/src/ext/ep/ht/kernels/internode.cu b/src/ext/ep/ht/kernels/internode.cu deleted file mode 100644 index b62ddf66..00000000 --- a/src/ext/ep/ht/kernels/internode.cu +++ /dev/null @@ -1,3011 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. -#include -#include -#include -#include - -#include "buffer.cuh" -#include "constants.cuh" -#include "device_helpers.cuh" -#include "exception.cuh" -#include "launch.cuh" - -namespace mscclpp { -namespace ep { - -// Packed type for loading NUM_MAX_NVL_PEERS bools as a single integer. -// Supports the two configurations the rest of the kernel logic assumes: -// 8 peers -> uint64_t (8 bools fit exactly) -// 4 peers -> uint32_t (4 bools fit exactly) -// All packed loads/stores below use this alias instead of a hardcoded uint64_t. -using NvlPackT = std::conditional_t; -static_assert(NUM_MAX_NVL_PEERS == 8 || NUM_MAX_NVL_PEERS == 4, - "NUM_MAX_NVL_PEERS must be 4 or 8 for HT internode kernel"); -static_assert(NUM_MAX_NVL_PEERS * sizeof(bool) == sizeof(NvlPackT), "NvlPackT size must match NUM_MAX_NVL_PEERS bools"); - -namespace internode { - -// =========================================================================== -// NVLS multimem PTX wrappers. `multimem.*` ops require sm_90+ (Hopper / -// Blackwell) and an active NVLink-SHARP multicast handle. Hosts gate -// every call site behind `nvls_mc_ptr != nullptr`, which is set only when -// `mscclpp::isNvlsSupported()` returned true at Buffer construction, so -// these macros are unreachable on older archs at runtime — the `__trap()` -// fallback is purely a build-time guard to keep PTX assembly clean on -// sm_80 / sm_86 / sm_89 targets. -// =========================================================================== -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) -#define MSCCLPP_EP_MULTIMEM_RED_ADD_RELAXED_U64(ptr, delta) \ - asm volatile("multimem.red.relaxed.sys.global.add.u64 [%0], %1;" ::"l"(ptr), "l"((uint64_t)(delta)) : "memory") -#define MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(ptr) \ - asm volatile("multimem.red.release.sys.global.add.u64 [%0], 1;" ::"l"(ptr) : "memory") -#define MSCCLPP_EP_MULTIMEM_ST_RELAXED_U32(ptr, val) \ - asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" ::"l"(ptr), "r"((uint32_t)(val)) : "memory") -#else -#define MSCCLPP_EP_MULTIMEM_RED_ADD_RELAXED_U64(ptr, delta) __trap() -#define MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(ptr) __trap() -#define MSCCLPP_EP_MULTIMEM_ST_RELAXED_U32(ptr, val) __trap() -#endif - -// =========================================================================== -// NVLS HT counter helpers (Phase 3). -// -// All cross-node tail/head counter updates that previously used HW -// atomicAdd (broken on Azure CX-7 RoCE: IBV_ATOMIC_NONE) and were -// patched to absolute-value RDMA WRITE — go through the NVLS multicast -// fabric instead. multimem.red.add is hardware-atomic across all ranks -// in the multicast group (validated cross-node Phase 2). For the legacy -// IB-platform path, callers pass `nvls_*_mc/dev == nullptr` and the -// existing PortChannel path runs. -// -// Layout (matches `nvls_off_tail/head` allocation in buffer.cc): -// region[kind][channel][src_rdma][dst_rdma] (uint64 each) -// Sender writes to slot(channel, my_rdma_rank, dst_rdma_rank). -// Receiver reads from slot(channel, src_rdma_rank, my_rdma_rank). -// Self-loop (src == dst == my_rdma_rank) goes through the same path. -// =========================================================================== - -__device__ __forceinline__ size_t nvls_ctr_slot_index(int channel_id, int src_rdma, int dst_rdma) { - // NUM_MAX_RDMA_PEERS² slots per channel; uint64 each → slot stride is 8B - // but we return slot index (multiply by sizeof(uint64) at use site via - // pointer arithmetic on uint64_t*). - return (size_t)channel_id * (size_t)NUM_MAX_RDMA_PEERS * (size_t)NUM_MAX_RDMA_PEERS + - (size_t)src_rdma * (size_t)NUM_MAX_RDMA_PEERS + (size_t)dst_rdma; -} - -__device__ __forceinline__ void nvls_ctr_add(void* mc_base, int channel_id, int src_rdma, int dst_rdma, - uint64_t delta) { - uint64_t* slot = reinterpret_cast(mc_base) + nvls_ctr_slot_index(channel_id, src_rdma, dst_rdma); - // NOTE: must be `relaxed`, not `release`. Empirically on Azure GB200, - // `multimem.red.release.sys.global.add.u64` issued concurrently from - // many channels (warps) triggers `unspecified launch failure`. Use a - // preceding `__threadfence_system()` at the call site to publish data - // before bumping the counter. - MSCCLPP_EP_MULTIMEM_RED_ADD_RELAXED_U64(slot, delta); -} - -__device__ __forceinline__ uint64_t nvls_ctr_load(void* dev_base, int channel_id, int src_rdma, int dst_rdma) { - uint64_t* slot = reinterpret_cast(dev_base) + nvls_ctr_slot_index(channel_id, src_rdma, dst_rdma); - uint64_t v; - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(slot) : "memory"); - return v; -} - -template -__global__ void __launch_bounds__(kNumThreads, 1) - get_dispatch_layout(const int64_t* topk_idx, int* num_tokens_per_rank, int* num_tokens_per_rdma_rank, - int* num_tokens_per_expert, bool* is_token_in_rank, int num_tokens, int num_topk, int num_ranks, - int num_experts) { - auto sm_id = static_cast(blockIdx.x); - auto thread_id = static_cast(threadIdx.x); - - // Count expert statistics - __shared__ int num_tokens_per_expert_per_thread[kNumThreads][kNumExpertsPerSM]; - int expert_begin_idx = sm_id * kNumExpertsPerSM, - expert_end_idx = min(expert_begin_idx + kNumExpertsPerSM, num_experts); - if (expert_begin_idx < expert_end_idx) { -// Per-thread count -#pragma unroll - for (int i = 0; i < kNumExpertsPerSM; ++i) num_tokens_per_expert_per_thread[thread_id][i] = 0; -#pragma unroll - for (int i = thread_id; i < num_tokens; i += kNumThreads) { - auto shifted_topk_idx = topk_idx + i * num_topk; -#pragma unroll - for (int j = 0, expert_idx; j < num_topk; ++j) { - expert_idx = static_cast(shifted_topk_idx[j]); - if (expert_begin_idx <= expert_idx and expert_idx < expert_end_idx) - ++num_tokens_per_expert_per_thread[thread_id][expert_idx - expert_begin_idx]; - } - } - __syncthreads(); - - // Sum up - EP_STATIC_ASSERT(kNumExpertsPerSM <= kNumThreads, "Too many experts per SM"); - if (expert_begin_idx + thread_id < expert_end_idx) { - int sum = 0; -#pragma unroll - for (int i = 0; i < kNumThreads; ++i) sum += num_tokens_per_expert_per_thread[i][thread_id]; - num_tokens_per_expert[expert_begin_idx + thread_id] = sum; - } - return; - } - - if (num_tokens_per_rdma_rank != nullptr) - EP_DEVICE_ASSERT(num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS); - - // Count rank statistics - constexpr int kNumRDMARanksPerSM = kNumRanksPerSM / NUM_MAX_NVL_PEERS; - __shared__ int num_tokens_per_rank_per_thread[kNumThreads][kNumRanksPerSM]; - __shared__ int num_tokens_per_rdma_rank_per_thread[kNumThreads][kNumRDMARanksPerSM]; - auto sm_begin = (num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM; - int rank_begin_idx = (sm_id - sm_begin) * kNumRanksPerSM, - rank_end_idx = min(rank_begin_idx + kNumRanksPerSM, num_ranks); - int rdma_rank_begin_idx = rank_begin_idx / NUM_MAX_NVL_PEERS, rdma_rank_end_idx = rank_end_idx / NUM_MAX_NVL_PEERS; - if (rank_begin_idx < rank_end_idx) { - const auto num_expert_per_rank = num_experts / num_ranks; - auto expert_begin = rank_begin_idx * num_expert_per_rank; - auto expert_end = rank_end_idx * num_expert_per_rank; - -// Per-thread count -#pragma unroll - for (int i = 0; i < kNumRanksPerSM; ++i) num_tokens_per_rank_per_thread[thread_id][i] = 0; -#pragma unroll - for (int i = 0; i < kNumRDMARanksPerSM; ++i) num_tokens_per_rdma_rank_per_thread[thread_id][i] = 0; -#pragma unroll - for (int i = thread_id; i < num_tokens; i += kNumThreads) { - auto shifted_topk_idx = topk_idx + i * num_topk; - int is_in_rank[kNumRanksPerSM] = {0}, is_in_rdma_rank[kNumRDMARanksPerSM] = {0}; -#pragma unroll - for (int j = 0, expert_idx, rank_idx; j < num_topk; ++j) { - expert_idx = static_cast(shifted_topk_idx[j]); - if (expert_begin <= expert_idx and expert_idx < expert_end) { - // Count single rank - rank_idx = expert_idx / num_expert_per_rank - rank_begin_idx; - is_in_rank[rank_idx]++, is_in_rdma_rank[rank_idx / NUM_MAX_NVL_PEERS]++; - } - } - - auto shifted_is_token_in_rank = is_token_in_rank + i * num_ranks; -#pragma unroll - for (int j = 0; j + rank_begin_idx < rank_end_idx; ++j) { - shifted_is_token_in_rank[j + rank_begin_idx] = (is_in_rank[j] > 0); - num_tokens_per_rank_per_thread[thread_id][j] += (is_in_rank[j] > 0); - } - -#pragma unroll - for (int j = 0; j + rdma_rank_begin_idx < rdma_rank_end_idx; ++j) - num_tokens_per_rdma_rank_per_thread[thread_id][j] += (is_in_rdma_rank[j] > 0); - } - __syncthreads(); - - // Sum up - EP_STATIC_ASSERT(kNumRanksPerSM <= kNumThreads, "Too many ranks per SM"); - if (rank_begin_idx + thread_id < rank_end_idx) { - int sum = 0; -#pragma unroll - for (int i = 0; i < kNumThreads; ++i) sum += num_tokens_per_rank_per_thread[i][thread_id]; - num_tokens_per_rank[rank_begin_idx + thread_id] = sum; - } - - if (num_tokens_per_rdma_rank != nullptr and rdma_rank_begin_idx + thread_id < rdma_rank_end_idx) { - int sum = 0; -#pragma unroll - for (int i = 0; i < kNumThreads; ++i) sum += num_tokens_per_rdma_rank_per_thread[i][thread_id]; - num_tokens_per_rdma_rank[rdma_rank_begin_idx + thread_id] = sum; - } - } -} - -void get_dispatch_layout(const int64_t* topk_idx, int* num_tokens_per_rank, int* num_tokens_per_rdma_rank, - int* num_tokens_per_expert, bool* is_token_in_rank, int num_tokens, int num_topk, - int num_ranks, int num_experts, cudaStream_t stream) { - constexpr int kNumThreads = 256, kNumExpertsPerSM = 32, kNumRanksPerSM = 8; - int num_sms = - ((num_experts + kNumExpertsPerSM - 1) / kNumExpertsPerSM) + (num_ranks + kNumRanksPerSM - 1) / kNumRanksPerSM; - EP_STATIC_ASSERT(kNumExpertsPerSM % NUM_MAX_NVL_PEERS == 0, "Invalid number of experts per SM"); - - SETUP_LAUNCH_CONFIG(num_sms, kNumThreads, stream); - LAUNCH_KERNEL(&cfg, (get_dispatch_layout), topk_idx, - num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, num_tokens, - num_topk, num_ranks, num_experts); -} - -struct SourceMeta { - int src_rdma_rank, is_token_in_nvl_rank_bits; - - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8 || NUM_MAX_NVL_PEERS == 4, "Invalid number of maximum NVL peers"); - - __forceinline__ SourceMeta() = default; - - // TODO: faster encoding - __device__ __forceinline__ SourceMeta(int rdma_rank, const bool* is_token_in_nvl_ranks) { - src_rdma_rank = rdma_rank; - is_token_in_nvl_rank_bits = is_token_in_nvl_ranks[0]; -#pragma unroll - for (int i = 1; i < NUM_MAX_NVL_PEERS; ++i) is_token_in_nvl_rank_bits |= is_token_in_nvl_ranks[i] << i; - } - - __device__ __forceinline__ bool is_token_in_nvl_rank(int nvl_rank) const { - return (is_token_in_nvl_rank_bits >> nvl_rank) & 1; - } -}; - -EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`"); - -int get_source_meta_bytes() { return sizeof(SourceMeta); } - -// Increment 6 (kEpFlat): post-dispatch metadata drain kernel. Under the flat path -// the sender wrote each token's metadata (SourceMeta + scales + topk_idx/weights, -// topk RAW) into the destination pool's META region at the token's final recv slot. -// This kernel runs on comm_stream right after the dispatch kernel (which, with the -// receiver still present, guarantees all senders' pool writes are visible by the -// time it exits), and copies the pool meta into the recv_* output tensors with the -// topk expert-range rebase the receiver would have done. One thread per recv token. -__global__ void flat_meta_drain_kernel(const uint8_t* __restrict__ pool_base, int64_t meta_base, int num_recv_tokens, - SourceMeta* __restrict__ recv_src_meta, float* __restrict__ recv_x_scales, - int64_t* __restrict__ recv_topk_idx, float* __restrict__ recv_topk_weights, - int num_scales, int num_topk, int expert_begin, int expert_end, - int64_t meta_slot_bytes) { - for (int t = blockIdx.x * blockDim.x + threadIdx.x; t < num_recv_tokens; t += gridDim.x * blockDim.x) { - const uint8_t* m = pool_base + meta_base + static_cast(t) * meta_slot_bytes; - if (recv_src_meta != nullptr) recv_src_meta[t] = *reinterpret_cast(m); - const float* sdat = reinterpret_cast(m + sizeof(SourceMeta)); - if (recv_x_scales != nullptr) - for (int s = 0; s < num_scales; ++s) recv_x_scales[static_cast(t) * num_scales + s] = sdat[s]; - const int* idat = reinterpret_cast(m + sizeof(SourceMeta) + num_scales * sizeof(float)); - const float* wdat = reinterpret_cast(idat + num_topk); - if (recv_topk_idx != nullptr) { - for (int k = 0; k < num_topk; ++k) { - const int idx = idat[k]; - const int local = (idx >= expert_begin and idx < expert_end) ? idx - expert_begin : -1; - recv_topk_idx[static_cast(t) * num_topk + k] = static_cast(local); - recv_topk_weights[static_cast(t) * num_topk + k] = (local >= 0) ? wdat[k] : 0.0f; - } - } - } -} - -void flat_meta_drain(void* pool_base, int64_t meta_base, int num_recv_tokens, void* recv_src_meta, float* recv_x_scales, - int64_t* recv_topk_idx, float* recv_topk_weights, int num_scales, int num_topk, int num_experts, - int num_ranks, int rank, int64_t meta_slot_bytes, cudaStream_t stream) { - if (num_recv_tokens <= 0) return; - const int num_local_experts = num_experts / num_ranks; - const int expert_begin = rank * num_local_experts; - const int expert_end = expert_begin + num_local_experts; - constexpr int kThreads = 256; - const int kBlocks = (num_recv_tokens + kThreads - 1) / kThreads; - flat_meta_drain_kernel<<>>( - reinterpret_cast(pool_base), meta_base, num_recv_tokens, - reinterpret_cast(recv_src_meta), recv_x_scales, recv_topk_idx, recv_topk_weights, num_scales, - num_topk, expert_begin, expert_end, meta_slot_bytes); -} - -__host__ __device__ __forceinline__ int get_num_bytes_per_rdma_token(int hidden_int4, int num_scales, int num_topk_idx, - int num_topk_weights) { - return static_cast(align(hidden_int4 * sizeof(int4) + sizeof(SourceMeta) + num_scales * sizeof(float) + - num_topk_idx * sizeof(int) + num_topk_weights * sizeof(float), - sizeof(int4))); -} - -__host__ __device__ __forceinline__ std::pair get_rdma_clean_meta(int hidden_int4, int num_scales, - int num_topk_idx, int num_topk_weights, - int num_rdma_ranks, - int num_rdma_recv_buffer_tokens, - int num_sms) { - // Return `int32_t` offset and count to clean - return {(get_num_bytes_per_rdma_token(hidden_int4, num_scales, num_topk_idx, num_topk_weights) * - num_rdma_recv_buffer_tokens * num_rdma_ranks * 2 * num_sms) / - sizeof(int), - (NUM_MAX_NVL_PEERS * 2 + 4) * num_rdma_ranks * 2 * num_sms}; -} - -__host__ __device__ __forceinline__ std::pair get_nvl_clean_meta(int hidden_int4, int num_scales, - int num_topk_idx, int num_topk_weights, - int num_rdma_ranks, int num_nvl_ranks, - int num_nvl_recv_buffer_tokens, - int num_sms) { - // Return `int32_t` offset and to clean - EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`"); - return { - (num_nvl_recv_buffer_tokens * - (hidden_int4 * sizeof(int4) + num_scales * sizeof(float) + num_topk_idx * sizeof(int) + - num_topk_weights * sizeof(float) + sizeof(SourceMeta)) * - num_nvl_ranks * num_sms) / - sizeof(int), - num_nvl_ranks * (2 * num_rdma_ranks + 2) * num_sms, - }; -} - -template -__global__ void notify_dispatch( - const int* num_tokens_per_rank, int* moe_recv_counter_mapped, int num_ranks, const int* num_tokens_per_rdma_rank, - int* moe_recv_rdma_counter_mapped, const int* num_tokens_per_expert, int* moe_recv_expert_counter_mapped, - int num_experts, const bool* is_token_in_rank, int num_tokens, int num_channels, int expert_alignment, - const int rdma_clean_offset, const int rdma_num_int_clean, const int nvl_clean_offset, const int nvl_num_int_clean, - int* rdma_channel_prefix_matrix, int* recv_rdma_rank_prefix_sum, int* gbl_channel_prefix_matrix, - int* recv_gbl_rank_prefix_sum, void* rdma_buffer_ptr, void** buffer_ptrs, int** task_fifo_ptrs, int head, int rank, - mscclpp::PortChannelDeviceHandle* port_channel_handles, mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, - // NVLS Phase 2 — replaces port_channel signal/wait + putWithSignal. - // When `nvls_mc_ptr == nullptr` the legacy PortChannel path runs - // unchanged (fallback for non-NVLS IB platforms). - void* nvls_mc_ptr, void* nvls_dev_ptr, size_t nvls_off_barrier, size_t nvls_off_data, uint64_t nvls_epoch, - int nvls_per_peer_bytes) { - auto sm_id = static_cast(blockIdx.x); - auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); - auto num_threads = static_cast(blockDim.x), num_warps = num_threads / 32; - - auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; - auto num_rdma_experts = num_experts / kNumRDMARanks, num_nvl_experts = num_rdma_experts / NUM_MAX_NVL_PEERS; - - if (sm_id == 0) { - // Communication with others - // Global barrier: the first warp do intra-node sync, the second warp do internode sync - EP_DEVICE_ASSERT(num_warps > 1); - EP_DEVICE_ASSERT(kNumRDMARanks + 32 <= num_threads); - const auto barrier_thread_id = thread_id - 32; - const bool run_barrier = - (barrier_thread_id >= 0) && (barrier_thread_id < kNumRDMARanks) && (barrier_thread_id != rdma_rank); - const auto barrier_channel_idx = - kLowLatencyMode ? barrier_thread_id : (barrier_thread_id * NUM_MAX_NVL_PEERS + nvl_rank); - if (nvls_mc_ptr != nullptr) { - // NVLS epoch barrier #0: every rank in the multicast group does - // `multimem.red.add 1` to slot[0] (over the NVL72 fabric — bypasses - // broken IB atomics on Azure CX-7 RoCE), then spins on its local - // mapping until the slot reaches `epoch * num_ranks` (everyone - // arrived). Replaces the pairwise PortChannel signal/wait above. - if (thread_id == 0) { - if (rank == 0) (void)0; - uint64_t* mc_b0 = reinterpret_cast(static_cast(nvls_mc_ptr) + nvls_off_barrier); - uint64_t* dev_b0 = reinterpret_cast(static_cast(nvls_dev_ptr) + nvls_off_barrier); - uint64_t pre; - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(pre) : "l"(dev_b0) : "memory"); - MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(mc_b0); - (void)0; - const uint64_t expected = nvls_epoch * static_cast(num_ranks); - uint64_t v; - int spin = 0; - do { - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(dev_b0) : "memory"); - if ((++spin & 0xFFFFFFF) == 0) { - (void)0; - } - } while (v < expected); - if (rank == 0) (void)0; - } - } else if (run_barrier) { - port_channel_handles[barrier_channel_idx].signal(); - port_channel_handles[barrier_channel_idx].wait(); - } - if constexpr (!kLowLatencyMode) { - // kLowLatencyMode==false requires sync of all ranks, which can be done by running intra-node sync - // after the inter-node sync is done. - __syncthreads(); - } -#if 1 - barrier_device(task_fifo_ptrs, head, nvl_rank); - move_fifo_slots(head); -#else - // TODO(chhwang): make memory channels work - if (thread_id < NUM_MAX_NVL_PEERS && thread_id != nvl_rank) { - memory_channel_handles[thread_id].relaxedSignal(); - memory_channel_handles[thread_id].relaxedWait(); - } -#endif - __syncthreads(); - - // Send numbers of tokens per rank/expert to RDMA ranks - auto rdma_buffer_ptr_int = reinterpret_cast(rdma_buffer_ptr); - auto num_elems = NUM_MAX_NVL_PEERS + num_rdma_experts + 1; - auto num_bytes = num_elems * sizeof(int); - auto per_channel_bytes = num_bytes * kNumRDMARanks; - auto rdma_recv_num_tokens_mixed = SymBuffer(rdma_buffer_ptr, num_elems, kNumRDMARanks); - - // Clean up for later data dispatch - EP_DEVICE_ASSERT(rdma_recv_num_tokens_mixed.total_bytes <= rdma_clean_offset * sizeof(int)); -#pragma unroll - for (int i = thread_id; i < rdma_num_int_clean; i += num_threads) rdma_buffer_ptr_int[rdma_clean_offset + i] = 0; - -// Copy to send buffer -#pragma unroll - for (int i = thread_id; i < num_ranks; i += num_threads) - rdma_recv_num_tokens_mixed.send_buffer(i / NUM_MAX_NVL_PEERS)[i % NUM_MAX_NVL_PEERS] = num_tokens_per_rank[i]; -#pragma unroll - for (int i = thread_id; i < num_experts; i += num_threads) - rdma_recv_num_tokens_mixed.send_buffer(i / num_rdma_experts)[NUM_MAX_NVL_PEERS + i % num_rdma_experts] = - num_tokens_per_expert[i]; - if (thread_id < kNumRDMARanks) - rdma_recv_num_tokens_mixed.send_buffer(thread_id)[NUM_MAX_NVL_PEERS + num_rdma_experts] = - num_tokens_per_rdma_rank[thread_id]; - __syncthreads(); - - // Issue send - // TODO: more light fence or barrier or signaling - // TODO: overlap EP barrier and NVL cleaning - if (nvls_mc_ptr != nullptr) { - // NVLS Phase 2 data path: every rank broadcasts its full send blob - // (kNumRDMARanks × num_bytes contiguous in send_buffer(0..)) into - // its own slot in the multicast data region using `multimem.st.u32`. - // After a second NVLS epoch barrier, every receiver copies the - // sub-block destined to it into the legacy `recv_buffer(s)` location - // so the rest of the kernel reads the same layout as before. - EP_DEVICE_ASSERT(num_bytes <= static_cast(nvls_per_peer_bytes)); - const int my_global_rank = kLowLatencyMode ? rdma_rank : (rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - const size_t slot_stride_bytes = - static_cast(NUM_MAX_RDMA_PEERS) * static_cast(nvls_per_peer_bytes); - // Sender: write our blob (num_elems × kNumRDMARanks ints) to our slot. - // SymBuffer::send_buffer(i) lives at base + i*num_elems ints in - // the local send region — contiguous from send_buffer(0). - int* src_ints = rdma_recv_num_tokens_mixed.send_buffer(0); - const int total_ints = num_elems * kNumRDMARanks; - int* mc_slot = reinterpret_cast(static_cast(nvls_mc_ptr) + nvls_off_data + - static_cast(my_global_rank) * slot_stride_bytes); - for (int i = thread_id; i < total_ints; i += num_threads) { - int val = src_ints[i]; - MSCCLPP_EP_MULTIMEM_ST_RELAXED_U32(mc_slot + i, val); - } - __syncthreads(); - // NVLS epoch barrier #1: ensure every sender's multimem.st has been - // delivered to all peers before we read. - if (thread_id == 0) { - if (rank == 0) (void)0; - uint64_t* mc_b1 = reinterpret_cast(static_cast(nvls_mc_ptr) + nvls_off_barrier + 8); - uint64_t* dev_b1 = reinterpret_cast(static_cast(nvls_dev_ptr) + nvls_off_barrier + 8); - MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(mc_b1); - const uint64_t expected = nvls_epoch * static_cast(num_ranks); - uint64_t v; - do { - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(dev_b1) : "memory"); - } while (v < expected); - if (rank == 0) (void)0; - } - __syncthreads(); - // Receiver: for each sender s, copy num_elems ints from the NVLS slot - // (sub-block at offset rdma_rank * num_bytes within the slot) into - // the legacy SymBuffer recv_buffer(s) location. - for (int s = 0; s < kNumRDMARanks; ++s) { - const int sender_global = kLowLatencyMode ? s : (s * NUM_MAX_NVL_PEERS + nvl_rank); - const int* nvls_src = reinterpret_cast(static_cast(nvls_dev_ptr) + nvls_off_data + - static_cast(sender_global) * slot_stride_bytes + - static_cast(rdma_rank) * num_bytes); - int* dst = rdma_recv_num_tokens_mixed.recv_buffer(s); - for (int i = thread_id; i < num_elems; i += num_threads) { - dst[i] = nvls_src[i]; - } - } - } else if (thread_id < kNumRDMARanks) { - auto dst_offset = rdma_rank * num_bytes + per_channel_bytes; - auto src_offset = thread_id * num_bytes; - auto peer_rank = kLowLatencyMode ? thread_id : (thread_id * NUM_MAX_NVL_PEERS + nvl_rank); - port_channel_handles[peer_rank].putWithSignal(dst_offset, src_offset, num_bytes); - port_channel_handles[peer_rank].wait(); - } - __syncthreads(); - - // NVL buffers - auto nvl_send_buffer = thread_id < NUM_MAX_NVL_PEERS ? buffer_ptrs[thread_id] : nullptr; - auto nvl_recv_buffer = buffer_ptrs[nvl_rank]; - auto nvl_reduced_num_tokens_per_expert = - Buffer(nvl_recv_buffer, num_rdma_experts).advance_also(nvl_send_buffer); - auto nvl_send_num_tokens_per_rank = AsymBuffer(nvl_send_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS); - auto nvl_send_num_tokens_per_expert = AsymBuffer(nvl_send_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS); - auto nvl_recv_num_tokens_per_rank = AsymBuffer(nvl_recv_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS); - auto nvl_recv_num_tokens_per_expert = AsymBuffer(nvl_recv_buffer, num_nvl_experts, NUM_MAX_NVL_PEERS); - - // Clean up for later data dispatch - auto nvl_buffer_ptr_int = reinterpret_cast(buffer_ptrs[nvl_rank]); - EP_DEVICE_ASSERT(nvl_reduced_num_tokens_per_expert.total_bytes + nvl_send_num_tokens_per_rank.total_bytes + - nvl_send_num_tokens_per_expert.total_bytes <= - nvl_clean_offset * sizeof(int)); -#pragma unroll - for (int i = thread_id; i < nvl_num_int_clean; i += num_threads) nvl_buffer_ptr_int[nvl_clean_offset + i] = 0; - - // Reduce number of tokens per expert into the NVL send buffer - // TODO: may use NVSHMEM reduction - EP_DEVICE_ASSERT(num_rdma_experts <= num_threads); - if (thread_id < num_rdma_experts) { - int sum = 0; -#pragma unroll - for (int i = 0; i < kNumRDMARanks; ++i) - sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + thread_id]; - nvl_reduced_num_tokens_per_expert[thread_id] = sum; - } - __syncthreads(); - - // Reduce RDMA received tokens - if (thread_id == 0) { - int sum = 0; -#pragma unroll - for (int i = 0; i < kNumRDMARanks; ++i) { - sum += rdma_recv_num_tokens_mixed.recv_buffer(i)[NUM_MAX_NVL_PEERS + num_rdma_experts]; - recv_rdma_rank_prefix_sum[i] = sum; - } - while (ld_volatile_global(moe_recv_rdma_counter_mapped) != -1); - *moe_recv_rdma_counter_mapped = sum; - } - - // Send numbers of tokens per rank/expert to NVL ranks - EP_DEVICE_ASSERT(NUM_MAX_NVL_PEERS <= num_threads); - if (thread_id < NUM_MAX_NVL_PEERS) { -#pragma unroll - for (int i = 0; i < kNumRDMARanks; ++i) - nvl_send_num_tokens_per_rank.buffer(nvl_rank)[i] = rdma_recv_num_tokens_mixed.recv_buffer(i)[thread_id]; -#pragma unroll - for (int i = 0; i < num_nvl_experts; ++i) - nvl_send_num_tokens_per_expert.buffer(nvl_rank)[i] = - nvl_reduced_num_tokens_per_expert[thread_id * num_nvl_experts + i]; - } - memory_fence(); - __syncthreads(); - barrier_device(task_fifo_ptrs, head, nvl_rank); - move_fifo_slots(head); - __syncthreads(); - - // Reduce number of tokens per rank/expert - EP_DEVICE_ASSERT(num_nvl_experts <= num_threads); - if (thread_id == 0) { - int sum = 0; -#pragma unroll - for (int i = 0; i < num_ranks; ++i) { - int src_rdma_rank = i / NUM_MAX_NVL_PEERS, src_nvl_rank = i % NUM_MAX_NVL_PEERS; - sum += nvl_recv_num_tokens_per_rank.buffer(src_nvl_rank)[src_rdma_rank]; - recv_gbl_rank_prefix_sum[i] = sum; - } - while (ld_volatile_global(moe_recv_counter_mapped) != -1); - *moe_recv_counter_mapped = sum; - } - if (thread_id < num_nvl_experts) { - int sum = 0; -#pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) sum += nvl_recv_num_tokens_per_expert.buffer(i)[thread_id]; - sum = (sum + expert_alignment - 1) / expert_alignment * expert_alignment; - while (ld_volatile_global(moe_recv_expert_counter_mapped + thread_id) != -1); - moe_recv_expert_counter_mapped[thread_id] = sum; - } - - // Finally barrier - __syncthreads(); - if (thread_id == 0) (void)0; - - if (nvls_mc_ptr != nullptr) { - // Phase 4 / NVLS HT "B2": replace the cross-node port_channel - // signal/wait pair with an epoch-monotonic NVLS multimem barrier - // (slot at +16). Bypasses the broken IB signal/wait path on Azure - // CX-7 RoCE that hangs the entire CTA at the !kLowLatencyMode - // __syncthreads below (thread 33's wait never completes). - if (thread_id == 0) { - if (rank == 0) (void)0; - uint64_t* mc_b2 = reinterpret_cast(static_cast(nvls_mc_ptr) + nvls_off_barrier + 16); - uint64_t* dev_b2 = reinterpret_cast(static_cast(nvls_dev_ptr) + nvls_off_barrier + 16); - MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(mc_b2); - const uint64_t expected = nvls_epoch * static_cast(num_ranks); - uint64_t v; - do { - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(dev_b2) : "memory"); - } while (v < expected); - if (rank == 0) (void)0; - } - } else if (run_barrier) { - port_channel_handles[barrier_channel_idx].signal(); - port_channel_handles[barrier_channel_idx].wait(); - } - if (thread_id == 0) (void)0; - if constexpr (!kLowLatencyMode) { - // kLowLatencyMode==false requires sync of all ranks, which can be done by running intra-node sync - // after the inter-node sync is done. - __syncthreads(); - } - barrier_device(task_fifo_ptrs, head, nvl_rank); - move_fifo_slots(head); - if (thread_id == 0) (void)0; - } else { - // Calculate meta data - int dst_rdma_rank = sm_id - 1; - for (int channel_id = warp_id; channel_id < num_channels; channel_id += num_warps) { - int token_start_idx, token_end_idx; - get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); - - // Iterate over tokens - int total_count = 0, per_nvl_rank_count[NUM_MAX_NVL_PEERS] = {0}; - for (int64_t i = token_start_idx + lane_id; i < token_end_idx; i += 32) { - auto is_token_in_rank_packed = - *reinterpret_cast(is_token_in_rank + i * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS); - auto is_token_in_rank_values = reinterpret_cast(&is_token_in_rank_packed); -#pragma unroll - for (int j = 0; j < NUM_MAX_NVL_PEERS; ++j) per_nvl_rank_count[j] += is_token_in_rank_values[j]; - total_count += (is_token_in_rank_packed != 0); - } - - // Warp reduce - total_count = warp_reduce_sum(total_count); -#pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) per_nvl_rank_count[i] = warp_reduce_sum(per_nvl_rank_count[i]); - - // Write into channel matrix - if (lane_id == 0) { -#pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) - gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + i) * num_channels + channel_id] = - per_nvl_rank_count[i]; - rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] = total_count; - } - } - - // Calculate prefix sum - __syncthreads(); - if (thread_id == 0) { - auto prefix_row = rdma_channel_prefix_matrix + dst_rdma_rank * num_channels; -#pragma unroll - for (int i = 1; i < num_channels; ++i) prefix_row[i] += prefix_row[i - 1]; - } - - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); - if (thread_id < NUM_MAX_NVL_PEERS) { - auto prefix_row = gbl_channel_prefix_matrix + (dst_rdma_rank * NUM_MAX_NVL_PEERS + thread_id) * num_channels; -#pragma unroll - for (int i = 1; i < num_channels; ++i) prefix_row[i] += prefix_row[i - 1]; - } - } -} - -void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, int num_ranks, - const int* num_tokens_per_rdma_rank, int* moe_recv_rdma_counter_mapped, - const int* num_tokens_per_expert, int* moe_recv_expert_counter_mapped, int num_experts, - const bool* is_token_in_rank, int num_tokens, int num_channels, int hidden_int4, int num_scales, - int num_topk, int expert_alignment, int* rdma_channel_prefix_matrix, - int* recv_rdma_rank_prefix_sum, int* gbl_channel_prefix_matrix, int* recv_gbl_rank_prefix_sum, - void* rdma_buffer_ptr, int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, - int num_max_nvl_chunked_recv_tokens, int** task_fifo_ptrs, int head, int rank, cudaStream_t stream, - int64_t num_rdma_bytes, int64_t num_nvl_bytes, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_mc_ptr, void* nvls_dev_ptr, - size_t nvls_off_barrier, size_t nvls_off_data, uint64_t nvls_epoch, int nvls_per_peer_bytes) { -#define NOTIFY_DISPATCH_LAUNCH_CASE(num_rdma_ranks) \ - { \ - auto notify_dispatch_func = \ - low_latency_mode ? notify_dispatch : notify_dispatch; \ - LAUNCH_KERNEL(&cfg, notify_dispatch_func, num_tokens_per_rank, moe_recv_counter_mapped, num_ranks, \ - num_tokens_per_rdma_rank, moe_recv_rdma_counter_mapped, num_tokens_per_expert, \ - moe_recv_expert_counter_mapped, num_experts, is_token_in_rank, num_tokens, num_channels, \ - expert_alignment, rdma_clean_meta.first, rdma_clean_meta.second, nvl_clean_meta.first, \ - nvl_clean_meta.second, rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, \ - gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, rdma_buffer_ptr, buffer_ptrs, task_fifo_ptrs, \ - head, rank, port_channel_handles, memory_channel_handles, nvls_mc_ptr, nvls_dev_ptr, \ - nvls_off_barrier, nvls_off_data, nvls_epoch, nvls_per_peer_bytes); \ - } \ - break - - constexpr int kNumThreads = 512; - const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; - - // Get clean meta. inc5: under MSCCLPP_EP_DIRECT the rdma ring slot excludes - // hidden (kEpDirect), so the clean region must match the SMALL slot the kernel - // uses, else the cleaned head/tail region != the one the kernel reads. - const bool ep_direct = []() { - const char* e = std::getenv("MSCCLPP_EP_DIRECT"); - return e && std::atoi(e) != 0; - }(); - auto rdma_clean_meta = get_rdma_clean_meta(ep_direct ? 0 : hidden_int4, num_scales, num_topk, num_topk, - num_rdma_ranks, num_max_rdma_chunked_recv_tokens, num_channels); - auto nvl_clean_meta = get_nvl_clean_meta(hidden_int4, num_scales, num_topk, num_topk, num_rdma_ranks, - NUM_MAX_NVL_PEERS, num_max_nvl_chunked_recv_tokens, num_channels); - EP_HOST_ASSERT((rdma_clean_meta.first + rdma_clean_meta.second) * sizeof(int) <= num_rdma_bytes); - EP_HOST_ASSERT((nvl_clean_meta.first + nvl_clean_meta.second) * sizeof(int) <= num_nvl_bytes); - EP_HOST_ASSERT(num_rdma_bytes < std::numeric_limits::max()); - EP_HOST_ASSERT(num_nvl_bytes < std::numeric_limits::max()); - - // Launch kernel - SETUP_LAUNCH_CONFIG(1 + num_rdma_ranks, kNumThreads, stream); - SWITCH_RDMA_RANKS(NOTIFY_DISPATCH_LAUNCH_CASE); -#undef NOTIFY_DISPATCH_LAUNCH_CASE -} - -// At most 8 RDMA ranks to be sent -constexpr int get_num_topk_rdma_ranks(int num_rdma_ranks) { return num_rdma_ranks < 8 ? num_rdma_ranks : 8; } - -template -__global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32), 1) - dispatch(int4* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv_topk_weights, - SourceMeta* recv_src_meta, const int4* x, const float* x_scales, const int64_t* topk_idx, - const float* topk_weights, int* send_rdma_head, int* send_nvl_head, int* recv_rdma_channel_prefix_matrix, - int* recv_gbl_channel_prefix_matrix, const int* rdma_channel_prefix_matrix, - const int* recv_rdma_rank_prefix_sum, const int* gbl_channel_prefix_matrix, - const int* recv_gbl_rank_prefix_sum, int num_tokens, int hidden_int4, int num_scales, int num_topk, - int num_experts, const bool* is_token_in_rank, void* rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens, - int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, int num_max_nvl_chunked_send_tokens, - int num_max_nvl_chunked_recv_tokens, int rank, int num_ranks, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, - // Phase 3 NVLS counter pointers (nullptr → fall back to PortChannel/atomicAdd path). - void* nvls_head_mc, void* nvls_head_dev, void* nvls_tail_mc, void* nvls_tail_dev, - // Phase 4: per-peer fabric-IPC base pointers; when non-null, cross-node data - // PUTs go directly through NVL72 fabric VA instead of `handle.put` over IB. - void* const* peer_rdma_bases) { - enum class WarpRole { - kRDMASender, - kRDMASenderCoordinator, - kRDMAAndNVLForwarder, - kForwarderCoordinator, - kNVLReceivers - }; - - const auto sm_id = static_cast(blockIdx.x); - const auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); - const auto num_channels = static_cast(gridDim.x) / 2, channel_id = sm_id / 2; - const bool is_forwarder = sm_id % 2 == 0; - const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; - - const auto role_meta = [=]() -> std::pair { - if (is_forwarder) { - if (warp_id < NUM_MAX_NVL_PEERS) { - return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % NUM_MAX_NVL_PEERS}; - } else { - return {WarpRole::kForwarderCoordinator, warp_id - NUM_MAX_NVL_PEERS}; - } - } else if (warp_id < kNumDispatchRDMASenderWarps) { - return {WarpRole::kRDMASender, -1}; - } else if (warp_id == kNumDispatchRDMASenderWarps) { - return {WarpRole::kRDMASenderCoordinator, -1}; - } else { - return {WarpRole::kNVLReceivers, (warp_id + channel_id - kNumDispatchRDMASenderWarps) % NUM_MAX_NVL_PEERS}; - } - }(); - auto warp_role = role_meta.first; - auto target_rank = role_meta.second; // Not applicable for RDMA senders - EP_DEVICE_ASSERT(num_warps == kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS); - if (thread_id == 0 && sm_id == 0) { - (void)0; - } - - // Data checks - EP_DEVICE_ASSERT(num_topk <= 32); - - // RDMA symmetric layout (packed-bool size guard is at namespace scope via NvlPackT). - // Snapshot the original base before SymBuffer constructors advance it; used - // below to compute MR-relative offsets for handle.put(). - void* const rdma_buffer_ptr_base = rdma_buffer_ptr; - auto hidden_bytes = hidden_int4 * sizeof(int4); - auto num_bytes_per_rdma_token = get_num_bytes_per_rdma_token(hidden_int4, num_scales, num_topk, num_topk); - auto rdma_channel_data = - SymBuffer(rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token, kNumRDMARanks, - channel_id, num_channels); - auto rdma_channel_meta = - SymBuffer(rdma_buffer_ptr, NUM_MAX_NVL_PEERS * 2 + 2, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_head = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_tail = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - // Scratch slots (one uint64_t per (channel, peer) per kind) holding the - // absolute counter values used as RDMA WRITE source. Replaces the broken - // HW atomicAdd path on Azure CX-7 RoCE (IBV_ATOMIC_NONE) — each tail/head - // slot has a single writer per peer, so atomicity is unnecessary; an - // absolute-value RDMA WRITE through the same QP as the data PUT preserves - // ordering by IB semantics. - auto rdma_channel_tail_send_src = - SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_head_send_src = - SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - - auto data_send_offset = - sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * kNumRDMARanks * channel_id; - auto data_recv_offset = sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * - kNumRDMARanks * (channel_id + num_channels); - auto meta_offset = - sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * kNumRDMARanks * num_channels * 2; - auto meta_send_offset = meta_offset + sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2) * kNumRDMARanks * channel_id; - auto meta_recv_offset = - meta_offset + sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2) * kNumRDMARanks * (channel_id + num_channels); - auto head_offset = meta_offset + sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2) * kNumRDMARanks * num_channels * 2; - auto head_send_offset = head_offset + sizeof(uint64_t) * kNumRDMARanks * channel_id; - auto tail_offset = head_offset + sizeof(uint64_t) * kNumRDMARanks * num_channels; - auto tail_send_offset = tail_offset + sizeof(uint64_t) * kNumRDMARanks * channel_id; - - // NVL buffer layouts - // NOTES: `rs_wr_buffer_ptr` means "Read for Senders, Write for Receivers", `ws_rr_buffer_ptr` means "Write for - // Senders, Read for Receivers" - void *rs_wr_buffer_ptr = nullptr, *ws_rr_buffer_ptr = nullptr; - int rs_wr_rank = 0, ws_rr_rank = 0; - if (warp_role == WarpRole::kRDMAAndNVLForwarder) - rs_wr_buffer_ptr = buffer_ptrs[nvl_rank], ws_rr_buffer_ptr = buffer_ptrs[target_rank], rs_wr_rank = nvl_rank, - ws_rr_rank = target_rank; - if (warp_role == WarpRole::kNVLReceivers) - rs_wr_buffer_ptr = buffer_ptrs[target_rank], ws_rr_buffer_ptr = buffer_ptrs[nvl_rank], rs_wr_rank = target_rank, - ws_rr_rank = nvl_rank; - - // Allocate buffers - auto nvl_channel_x = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * hidden_int4, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_src_meta = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_x_scales = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_scales, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_topk_idx = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_topk_weights = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_prefix_start = - AsymBuffer(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_prefix_end = - AsymBuffer(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_head = AsymBuffer(rs_wr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id, num_channels, ws_rr_rank) - .advance_also(ws_rr_buffer_ptr); - auto nvl_channel_tail = AsymBuffer(ws_rr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - - // RDMA sender warp synchronization - __shared__ volatile int rdma_send_next_token_idx; - __shared__ volatile int rdma_send_channel_tail[kNumRDMARanks]; - __shared__ volatile int rdma_send_channel_next_tail[kNumRDMARanks]; - auto sync_rdma_sender_smem = []() { asm volatile("bar.sync 0, %0;" ::"r"((kNumDispatchRDMASenderWarps + 1) * 32)); }; - - // Forward warp synchronization - __shared__ volatile int forward_channel_head[NUM_MAX_NVL_PEERS][kNumRDMARanks]; - __shared__ volatile bool forward_channel_retired[NUM_MAX_NVL_PEERS]; - auto sync_forwarder_smem = []() { asm volatile("bar.sync 1, %0;" ::"r"((NUM_MAX_NVL_PEERS + 1) * 32)); }; - - if (warp_role == WarpRole::kRDMASender) { - // Get tasks - int token_start_idx, token_end_idx; - get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); - - // Clean shared memory - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA ranks"); - (warp_id == 0 and lane_id == 0) ? (rdma_send_next_token_idx = token_start_idx) : 0; - (warp_id == 0 and lane_id < kNumRDMARanks) ? (rdma_send_channel_tail[lane_id] = 0) : 0; - (warp_id == 0 and lane_id < kNumRDMARanks) ? (rdma_send_channel_next_tail[lane_id] = 0) : 0; - - // Send number of tokens in this channel by `-value - 1` - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * 2 + 2 <= 32, "Invalid number of NVL peers"); - for (int dst_rdma_rank = warp_id; dst_rdma_rank < kNumRDMARanks; dst_rdma_rank += kNumDispatchRDMASenderWarps) { - auto dst_ptr = dst_rdma_rank == rdma_rank ? rdma_channel_meta.recv_buffer(dst_rdma_rank) - : rdma_channel_meta.send_buffer(dst_rdma_rank); - if (lane_id < NUM_MAX_NVL_PEERS) { - dst_ptr[lane_id] = - -(channel_id == 0 ? 0 - : gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id) * num_channels + - channel_id - 1]) - - 1; - } else if (lane_id < NUM_MAX_NVL_PEERS * 2) { - dst_ptr[lane_id] = - -gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id - NUM_MAX_NVL_PEERS) * - num_channels + - channel_id] - - 1; - } else if (lane_id == NUM_MAX_NVL_PEERS * 2) { - dst_ptr[lane_id] = - -(channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1]) - 1; - } else if (lane_id == NUM_MAX_NVL_PEERS * 2 + 1) { - dst_ptr[lane_id] = -rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] - 1; - } - __syncwarp(); - - if (dst_rdma_rank == rdma_rank) continue; - - // Issue RDMA for non-local ranks - if (peer_rdma_bases != nullptr) { - // Phase 4: deliver meta directly via NVL72 fabric VA (cooperative - // int copy across the warp). Replaces port_channel.put which is - // unreliable cross-node on Azure CX-7 RoCE without flush. - const auto num_bytes = sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2); - const auto dst_offset = rdma_rank * num_bytes + meta_recv_offset; - const auto src_offset = dst_rdma_rank * num_bytes + meta_send_offset; - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const int* src_p = reinterpret_cast(reinterpret_cast(rdma_buffer_ptr_base) + src_offset); - int* dst_p = reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + dst_offset); - const int n_int = (int)(num_bytes / sizeof(int)); - for (int k = lane_id; k < n_int; k += 32) { - dst_p[k] = src_p[k]; - } - __threadfence_system(); - } else if (lane_id == 0) { - auto num_bytes = sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2); - auto dst_offset = rdma_rank * num_bytes + meta_recv_offset; - auto src_offset = dst_rdma_rank * num_bytes + meta_send_offset; - auto port_channel_idx = kLowLatencyMode - ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - port_channel_handles[port_channel_idx].put(dst_offset, src_offset, num_bytes); - // port_channel_handles[port_channel_idx].flush(); - } - __syncwarp(); - } - sync_rdma_sender_smem(); - - // Iterate over tokens and copy into buffer - int64_t token_idx; - int cached_rdma_channel_head = 0, last_rdma_tail_idx = -1; - auto send_buffer = - lane_id == rdma_rank ? rdma_channel_data.recv_buffer(lane_id) : rdma_channel_data.send_buffer(lane_id); - for (token_idx = token_start_idx + warp_id; token_idx < token_end_idx; token_idx += kNumDispatchRDMASenderWarps) { - // Read RDMA rank existence - NvlPackT is_token_in_rank_uint64 = 0; - if (lane_id < kNumRDMARanks) - is_token_in_rank_uint64 = - *reinterpret_cast(is_token_in_rank + token_idx * num_ranks + lane_id * NUM_MAX_NVL_PEERS); - - // Acquire sequential lock - while (lane_id == 0 and rdma_send_next_token_idx != token_idx); - __syncwarp(); - - // Acquire next tail - int rdma_tail_idx = -1; - if (is_token_in_rank_uint64 != 0) { - rdma_tail_idx = rdma_send_channel_next_tail[lane_id]++; - while (rdma_tail_idx - cached_rdma_channel_head >= num_max_rdma_chunked_recv_tokens) { - // Phase 4: head feedback path \u2014 cross-node uses fabric-VA store, - // self-loop uses local atomic. Both end up in rdma_channel_head; - // single read via ld_volatile_global covers them. NVLS removed. - cached_rdma_channel_head = static_cast(ld_volatile_global(rdma_channel_head.buffer(lane_id))); - } - } - __syncwarp(); - - // Store RDMA head for combine - if (lane_id < kNumRDMARanks and not kCachedMode) - send_rdma_head[token_idx * kNumRDMARanks + lane_id] = rdma_tail_idx; - - // Update last token tail. In-loop writes are sequenced by the - // per-channel sequential lock and the warp-stride property of the - // token loop, so monotonicity is guaranteed and a plain - // st_release_cta is correct AND faster than atomicMax (which - // would serialize through L2 if the compiler can't infer shared - // address space). The epilogue (out of the seq-lock contract for - // the highest in-rank slot) needs atomicMax separately. - if (last_rdma_tail_idx >= 0) - st_release_cta(const_cast(rdma_send_channel_tail + lane_id), last_rdma_tail_idx + 1); - last_rdma_tail_idx = rdma_tail_idx; - - // Release sequential lock - lane_id == 0 ? (rdma_send_next_token_idx += 1) : 0; - - // Broadcast tails - SourceMeta src_meta; - int num_topk_ranks = 0, topk_ranks[kNumTopkRDMARanks]; - void* dst_send_buffers[kNumTopkRDMARanks]; -#pragma unroll - for (int i = 0, slot_idx; i < kNumRDMARanks; ++i) - if ((slot_idx = __shfl_sync(0xffffffff, rdma_tail_idx, i)) >= 0) { - slot_idx = slot_idx % num_max_rdma_chunked_recv_tokens; - topk_ranks[num_topk_ranks] = i; - auto recv_is_token_in_rank_uint64 = warpBroadcast(is_token_in_rank_uint64, i); - auto recv_is_token_in_rank_values = reinterpret_cast(&recv_is_token_in_rank_uint64); - if (lane_id == num_topk_ranks) src_meta = SourceMeta(rdma_rank, recv_is_token_in_rank_values); - dst_send_buffers[num_topk_ranks++] = - reinterpret_cast(warpBroadcast(send_buffer, i)) + slot_idx * num_bytes_per_rdma_token; - } - EP_DEVICE_ASSERT(num_topk_ranks <= kNumTopkRDMARanks); - - // Copy `x` into symmetric send buffer - auto st_broadcast = [=](const int key, const int4& value) { -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) - st_na_global(reinterpret_cast(dst_send_buffers[j]) + key, value); - }; - UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast); -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) - dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + hidden_int4; - - // Copy source metadata into symmetric send buffer - if (lane_id < num_topk_ranks) st_na_global(reinterpret_cast(dst_send_buffers[lane_id]), src_meta); -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) - dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + 1; - -// Copy `x_scales` into symmetric send buffer -#pragma unroll - for (int i = lane_id; i < num_scales; i += 32) { - auto value = ld_nc_global(x_scales + token_idx * num_scales + i); -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) st_na_global(reinterpret_cast(dst_send_buffers[j]) + i, value); - } -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) - dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + num_scales; - -// Copy `topk_idx` and `topk_weights` into symmetric send buffer -#pragma unroll - for (int i = lane_id; i < num_topk * num_topk_ranks; i += 32) { - auto rank_idx = i / num_topk, copy_idx = i % num_topk; - auto idx_value = static_cast(ld_nc_global(topk_idx + token_idx * num_topk + copy_idx)); - auto weight_value = ld_nc_global(topk_weights + token_idx * num_topk + copy_idx); - st_na_global(reinterpret_cast(dst_send_buffers[rank_idx]) + copy_idx, idx_value); - st_na_global(reinterpret_cast(dst_send_buffers[rank_idx]) + num_topk + copy_idx, weight_value); - } - } - - // Epilogue - // Acquire sequential lock - while (lane_id == 0 and rdma_send_next_token_idx != token_idx); - __syncwarp(); - - // Update last token tail (epilogue). See in-loop note on atomicMax. - if (last_rdma_tail_idx >= 0) atomicMax(const_cast(rdma_send_channel_tail + lane_id), last_rdma_tail_idx + 1); - - // Release sequential lock - lane_id == 0 ? (rdma_send_next_token_idx += 1) : 0; - } else if (warp_role == WarpRole::kRDMASenderCoordinator) { - // NOTES: in case of splitting the issued put at the end of the buffer - EP_DEVICE_ASSERT(num_max_rdma_chunked_recv_tokens % num_max_rdma_chunked_send_tokens == 0); - - // Synchronize shared memory - sync_rdma_sender_smem(); - if (lane_id == 0 && channel_id == 0 && rank == 0) { - (void)0; - } - - // Get number of tokens to send for each RDMA rank - int num_tokens_to_send = 0; - if (lane_id < kNumRDMARanks) { - num_tokens_to_send = rdma_channel_prefix_matrix[lane_id * num_channels + channel_id]; - if (channel_id > 0) num_tokens_to_send -= rdma_channel_prefix_matrix[lane_id * num_channels + channel_id - 1]; - } - - // Iterate all RDMA ranks - int last_issued_tail = 0; - while (__any_sync(0xffffffff, num_tokens_to_send > 0)) { -#pragma unroll - for (int i = 0; i < kNumRDMARanks; ++i, __syncwarp()) { - // To mitigate incast congestion, shuffle the starting index of target rank for different ranks and channels - const int dst_rdma_rank = (i + channel_id + rdma_rank) % kNumRDMARanks; - - // ----------------------------------------------------------------- - // Phase 4 restructure: owner lane (lane_id == dst_rdma_rank) decides - // whether to issue, then broadcasts num_tokens_to_issue and the - // tail-base for slot indexing to all 32 lanes via shfl. All lanes - // then participate in the cooperative cross-node fabric-IPC copy - // (when peer_rdma_bases != nullptr). This replaces the legacy - // single-lane handle.put(data) path which is broken on Azure CX-7. - // ----------------------------------------------------------------- - const bool owner = (lane_id == dst_rdma_rank); - int my_num_tokens_to_issue = 0; - int my_issue_tail = last_issued_tail; - if (owner && num_tokens_to_send > 0) { - auto processed_tail = ld_acquire_cta(const_cast(rdma_send_channel_tail + dst_rdma_rank)); - auto num_tokens_processed = processed_tail - last_issued_tail; - if (num_tokens_processed == num_tokens_to_send || num_tokens_processed >= num_max_rdma_chunked_send_tokens) { - int n = min(num_tokens_processed, num_max_rdma_chunked_send_tokens); - EP_DEVICE_ASSERT(n >= 0 && n <= num_tokens_to_send); - my_num_tokens_to_issue = n; - } - } - - const int n_issue = __shfl_sync(0xffffffff, my_num_tokens_to_issue, dst_rdma_rank); - const int issue_tail = __shfl_sync(0xffffffff, my_issue_tail, dst_rdma_rank); - if (n_issue == 0) continue; - - if (nvls_tail_mc != nullptr) { - // Phase 3+4: NVLS counter fast path. For cross-node, payload - // delivery uses warp-cooperative direct fabric-IPC writes when - // peer_rdma_bases is available; otherwise falls back to single- - // lane handle.put + flush. - if (dst_rdma_rank != rdma_rank) { - const auto dst_slot_idx = issue_tail % num_max_rdma_chunked_recv_tokens; - const size_t num_bytes_per_msg = (size_t)num_bytes_per_rdma_token * (size_t)n_issue; - const auto dst_offset = rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_recv_offset; - const auto src_offset = dst_rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_send_offset; - if (peer_rdma_bases != nullptr) { - // Phase 4: warp-cooperative int4 stores via NVL72 fabric VA. - // Uses NVLS counter (nvls_ctr_add) below to publish the - // writes to the receiver — `multimem.red.RELAXED` + a - // preceding `__threadfence_system()` is the validated - // ordering pair (release semantics on multimem.red triggers - // unspecified launch failure on Azure GB200). - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const int4* src_p = - reinterpret_cast(reinterpret_cast(rdma_buffer_ptr_base) + src_offset); - int4* dst_p = - reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + dst_offset); - const int n_int4 = (int)(num_bytes_per_msg / sizeof(int4)); - // Unrolled 8x to give the LSU pipeline more outstanding stores - // per lane. Each lane handles k, k+32, ..., k+224 per iter - // (stride 256) before looping. Tail handled with stride-32 loop. - const int stride8 = 8 * 32; - int k = lane_id; - const int n_full = (n_int4 / stride8) * stride8; - for (; k < n_full; k += stride8) { - int4 v0 = src_p[k]; - int4 v1 = src_p[k + 32]; - int4 v2 = src_p[k + 64]; - int4 v3 = src_p[k + 96]; - int4 v4 = src_p[k + 128]; - int4 v5 = src_p[k + 160]; - int4 v6 = src_p[k + 192]; - int4 v7 = src_p[k + 224]; - dst_p[k] = v0; - dst_p[k + 32] = v1; - dst_p[k + 64] = v2; - dst_p[k + 96] = v3; - dst_p[k + 128] = v4; - dst_p[k + 160] = v5; - dst_p[k + 192] = v6; - dst_p[k + 224] = v7; - } - for (; k < n_int4; k += 32) { - dst_p[k] = src_p[k]; - } - __syncwarp(); - __threadfence_system(); - } else if (owner) { - const auto port_channel_idx = - kLowLatencyMode ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - handle.put(dst_offset, src_offset, num_bytes_per_msg); - handle.flush(); - } - } - // Owner advances tail counter for this peer. - // Phase 4 fix: cross-node tail goes through direct fabric-VA - // store on peer's rdma_channel_tail slot. Self-loop tail goes - // through plain local atomicAdd on rdma_channel_tail.buffer(rdma_rank) - // — NVLS multicast is WRONG for self-loop because it fans out - // to all bound NVL peers' buffers (4 NVL ranks × n_issue ⇒ 4x - // over-count on each consumer's read). - if (owner) { - if (peer_rdma_bases != nullptr && dst_rdma_rank != rdma_rank) { - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const uintptr_t my_tail_off = reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - uint64_t* peer_tail = reinterpret_cast( - reinterpret_cast(peer_rdma_bases[dst_rank_global]) + my_tail_off); - const uint64_t new_tail = (uint64_t)issue_tail + (uint64_t)n_issue; - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(peer_tail), "l"(new_tail) : "memory"); - } else { - // Self-loop: plain release atomic on local slot (no multicast). - mscclpp::atomicFetchAdd(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)), - (uint64_t)n_issue, mscclpp::memoryOrderRelease); - } - } - } else if (owner) { - // Legacy non-NVLS path (single-lane). - if (dst_rdma_rank == rdma_rank) { - // Update tails - mscclpp::atomicFetchAdd(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)), (uint64_t)n_issue, - mscclpp::memoryOrderRelease); - } else { - const auto dst_slot_idx = issue_tail % num_max_rdma_chunked_recv_tokens; - const size_t num_bytes_per_msg = (size_t)num_bytes_per_rdma_token * (size_t)n_issue; - const auto dst_offset = rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_recv_offset; - const auto src_offset = dst_rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_send_offset; - const auto port_channel_idx = kLowLatencyMode - ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - handle.put(dst_offset, src_offset, num_bytes_per_msg); - - // HW atomicAdd is broken on Azure CX-7 RoCE (IBV_ATOMIC_NONE). - // Write the new absolute tail to a local scratch slot, then RDMA - // WRITE it to the peer's tail recv slot. - const uint64_t new_tail = (uint64_t)issue_tail + (uint64_t)n_issue; - *rdma_channel_tail_send_src.buffer(dst_rdma_rank) = new_tail; - __threadfence_system(); - const auto src_off_tail = reinterpret_cast(rdma_channel_tail_send_src.buffer(dst_rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - handle.put(rdma_rank * sizeof(uint64_t) + tail_send_offset, src_off_tail, sizeof(uint64_t)); - } - } - if (owner) { - last_issued_tail += n_issue; - num_tokens_to_send -= n_issue; - } - } - } - // Phase 4 diag: report final issued tail per (channel, dst_rdma_rank). - if (lane_id < kNumRDMARanks && rank == 0) { - (void)0; - } - } else if (warp_role == WarpRole::kRDMAAndNVLForwarder) { - // RDMA consumers and NVL producers - const auto dst_nvl_rank = target_rank; - const auto dst_rank = rdma_rank * NUM_MAX_NVL_PEERS + dst_nvl_rank; - const auto dst_rank_expert_begin = dst_rank * (num_experts / num_ranks); - const auto dst_rank_expert_end = dst_rank_expert_begin + (num_experts / num_ranks); - - // Wait counters to arrive - int num_tokens_to_recv_from_rdma = 0, src_rdma_channel_prefix = 0; - EP_DEVICE_ASSERT(kNumRDMARanks <= 32); - auto start_time = clock64(); - if (lane_id < kNumRDMARanks) { - while (true) { - auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank); - auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS + dst_nvl_rank); - auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2); - auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2 + 1); - if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) { - // Notify NVL ranks - int start_sum = -meta_0 - 1, end_sum = -meta_1 - 1; - EP_DEVICE_ASSERT(start_sum >= 0 and end_sum >= 0 and end_sum >= start_sum); - st_relaxed_sys_global(nvl_channel_prefix_start.buffer() + lane_id, -start_sum - 1); - st_relaxed_sys_global(nvl_channel_prefix_end.buffer() + lane_id, -end_sum - 1); - - // Save RDMA channel received token count - src_rdma_channel_prefix = -meta_2 - 1; - auto src_rdma_channel_prefix_1 = -meta_3 - 1; - num_tokens_to_recv_from_rdma = src_rdma_channel_prefix_1 - src_rdma_channel_prefix; - if (not kCachedMode) - recv_rdma_channel_prefix_matrix[lane_id * num_channels + channel_id] = src_rdma_channel_prefix_1; - src_rdma_channel_prefix += lane_id == 0 ? 0 : recv_rdma_rank_prefix_sum[lane_id - 1]; - EP_DEVICE_ASSERT(num_tokens_to_recv_from_rdma >= 0); - // Phase 4 diag: report received expected token count per (channel, src). - if (rank == 4) { - (void)0; - } - break; - } - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch forwarder timeout (RDMA meta), channel: %d, RDMA: %d, nvl: %d, src RDMA lane: %d, dst " - "NVL: %d, meta: %d, %d, %d, %d\n", - channel_id, rdma_rank, nvl_rank, lane_id, dst_nvl_rank, meta_0, meta_1, meta_2, meta_3); - trap(); - } - } - } - __syncwarp(); - - // Shift cached head - send_nvl_head += src_rdma_channel_prefix * NUM_MAX_NVL_PEERS + dst_nvl_rank; - - // Wait shared memory to be cleaned - sync_forwarder_smem(); - - // Forward tokens from RDMA buffer - // NOTES: always start from the local rank - int src_rdma_rank = sm_id % kNumRDMARanks; - int cached_rdma_channel_head = 0, cached_rdma_channel_tail = 0; - int cached_nvl_channel_head = 0, cached_nvl_channel_tail = 0, rdma_nvl_token_idx = 0; - while (__any_sync(0xffffffff, num_tokens_to_recv_from_rdma > 0)) { - // Check destination queue emptiness, or wait a buffer to be released - start_time = clock64(); - while (lane_id == 0) { - int num_used_slots = cached_nvl_channel_tail - cached_nvl_channel_head; - if (num_max_nvl_chunked_recv_tokens - num_used_slots >= num_max_nvl_chunked_send_tokens) break; - cached_nvl_channel_head = ld_volatile_global(nvl_channel_head.buffer()); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch forwarder timeout (NVL check), channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, head: %d, " - "tail: %d\n", - channel_id, rdma_rank, nvl_rank, dst_nvl_rank, ld_volatile_global(nvl_channel_head.buffer()), - cached_nvl_channel_tail); - trap(); - } - } - __syncwarp(); - - // Find next source RDMA rank (round-robin) - start_time = clock64(); - while (true) { - src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks; - if (__shfl_sync(0xffffffff, num_tokens_to_recv_from_rdma, src_rdma_rank) > 0) { - if (lane_id == src_rdma_rank and cached_rdma_channel_head == cached_rdma_channel_tail) { - // Phase 4 fix: cross-node tail is now delivered via direct - // fabric-VA store to local rdma_channel_tail.buffer(src_rdma_rank), - // so prefer the legacy ld_acquire read which sees those stores. - // NVLS counter only used for self path (single rank). - // Phase 4 fix: cross-node tail comes via direct fabric-VA store - // (sender writes peer's rdma_channel_tail slot). Self-loop tail - // is a plain local atomic. Both end up in rdma_channel_tail — - // single read path via ld_acquire_sys_global covers them. - cached_rdma_channel_tail = static_cast(ld_acquire_sys_global(rdma_channel_tail.buffer(src_rdma_rank))); - } - if (__shfl_sync(0xffffffff, cached_rdma_channel_tail > cached_rdma_channel_head, src_rdma_rank)) break; - } - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { - printf( - "DeepEP dispatch forwarder timeout (RDMA check), channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, src RDMA " - "lane: %d, head: %d, tail: %d, expected: %d\n", - channel_id, rdma_rank, nvl_rank, dst_nvl_rank, lane_id, cached_rdma_channel_head, - cached_rdma_channel_tail, num_tokens_to_recv_from_rdma); - trap(); - } - } - auto src_rdma_head = __shfl_sync(0xffffffff, cached_rdma_channel_head, src_rdma_rank); - auto src_rdma_tail = __shfl_sync(0xffffffff, cached_rdma_channel_tail, src_rdma_rank); - - if (rank == 4 && lane_id == 0 && dst_nvl_rank == 0 && channel_id == 0) { - (void)0; - } - - // Iterate over every token from the RDMA buffer - for (int i = src_rdma_head, num_tokens_sent = 0; i < src_rdma_tail; ++i) { - auto rdma_slot_idx = i % num_max_rdma_chunked_recv_tokens; - void* shifted = rdma_channel_data.recv_buffer(src_rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token; - auto src_meta = ld_nc_global(reinterpret_cast(reinterpret_cast(shifted) + hidden_bytes)); - lane_id == src_rdma_rank ? (num_tokens_to_recv_from_rdma -= 1) : 0; - bool is_in_dst_nvl_rank = src_meta.is_token_in_nvl_rank(dst_nvl_rank); - if (lane_id == src_rdma_rank) { - auto cached_head = is_in_dst_nvl_rank ? rdma_nvl_token_idx : -1; - rdma_nvl_token_idx += is_in_dst_nvl_rank; - if (not kCachedMode) send_nvl_head[i * NUM_MAX_NVL_PEERS] = cached_head; - } - if (not is_in_dst_nvl_rank) continue; - - // Get an empty slot - int dst_slot_idx = (cached_nvl_channel_tail++) % num_max_nvl_chunked_recv_tokens; - - // Copy data - UNROLLED_WARP_COPY(5, lane_id, hidden_int4, nvl_channel_x.buffer() + dst_slot_idx * hidden_int4, - reinterpret_cast(shifted), ld_nc_global, st_na_global); - shifted = reinterpret_cast(shifted) + hidden_int4; - - // Copy source meta - if (lane_id == 0) st_na_global(nvl_channel_src_meta.buffer() + dst_slot_idx, src_meta); - shifted = reinterpret_cast(shifted) + 1; - - // Copy `x_scales` - UNROLLED_WARP_COPY(1, lane_id, num_scales, nvl_channel_x_scales.buffer() + dst_slot_idx * num_scales, - reinterpret_cast(shifted), ld_nc_global, st_na_global); - shifted = reinterpret_cast(shifted) + num_scales; - - // Copy `topk_idx` and `topk_weights` - // NOTES: do not use `shifted` after this `if`, because only several lanes are shifted - if (lane_id < num_topk) { - // Read - auto idx_value = ld_nc_global(reinterpret_cast(shifted) + lane_id); - shifted = reinterpret_cast(shifted) + num_topk; - auto weight_value = ld_nc_global(reinterpret_cast(shifted) + lane_id); - - // Transform and write - idx_value = (idx_value >= dst_rank_expert_begin and idx_value < dst_rank_expert_end) - ? idx_value - dst_rank_expert_begin - : -1; - st_na_global(nvl_channel_topk_idx.buffer() + dst_slot_idx * num_topk + lane_id, idx_value); - weight_value = idx_value >= 0 ? weight_value : 0.0f; - st_na_global(nvl_channel_topk_weights.buffer() + dst_slot_idx * num_topk + lane_id, weight_value); - } - - // In case of insufficient NVL buffers, early stopping - if ((++num_tokens_sent) == num_max_nvl_chunked_send_tokens) src_rdma_tail = i + 1; - } - - // Sync head index - if (lane_id == src_rdma_rank) - forward_channel_head[dst_nvl_rank][src_rdma_rank] = (cached_rdma_channel_head = src_rdma_tail); - - // Move tail index - __syncwarp(); - if (lane_id == 0) st_release_sys_global(nvl_channel_tail.buffer(), cached_nvl_channel_tail); - } - - // Retired - __syncwarp(); - if (lane_id == 0) { - forward_channel_retired[dst_nvl_rank] = true; - if (channel_id == 0) { - (void)0; - } - } - } else if (warp_role == WarpRole::kForwarderCoordinator) { - // Extra warps for forwarder coordinator should exit directly - if (target_rank > 0) return; - - // Forward warp coordinator - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - - // Clean shared memory - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); -#pragma unroll - for (int i = lane_id; i < kNumRDMARanks * NUM_MAX_NVL_PEERS; i += 32) - forward_channel_head[i % NUM_MAX_NVL_PEERS][i / NUM_MAX_NVL_PEERS] = 0; - if (lane_id < NUM_MAX_NVL_PEERS) forward_channel_retired[lane_id] = false; - sync_forwarder_smem(); - - int last_head = 0, target_rdma = lane_id < kNumRDMARanks ? lane_id : 0; - while (true) { - // Find minimum head - int min_head = std::numeric_limits::max(); -#pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) - if (not forward_channel_retired[i]) min_head = min(min_head, forward_channel_head[i][target_rdma]); - if (__all_sync(0xffffffff, min_head == std::numeric_limits::max())) break; - - // Update remote head - // Phase 4 perf: lower the lazy-update threshold from chunk_send → chunk_send/4 - // (or 1 if retired) so the sender's receive-buffer-space window - // advances more frequently. The original threshold caused partial - // last chunks to never trigger head feedback, which deadlocked at - // larger chunk_send values where 4096 tokens / 10 channels / 2 peers - // ≈ 205 tokens ⇒ 3 full chunks + 1 partial. - const bool any_retired = forward_channel_retired[0] || forward_channel_retired[1] || forward_channel_retired[2] || - forward_channel_retired[3] || forward_channel_retired[4] || forward_channel_retired[5] || - forward_channel_retired[6] || forward_channel_retired[7]; - const int head_update_threshold = any_retired ? 1 : max(1, num_max_rdma_chunked_send_tokens / 4); - if (min_head != std::numeric_limits::max() and min_head >= last_head + head_update_threshold and - lane_id < kNumRDMARanks) { - if (peer_rdma_bases != nullptr && lane_id != rdma_rank) { - // Phase 4 fix: cross-node head feedback via direct fabric-VA - // store on peer's rdma_channel_head slot (single writer per - // (channel, my_rdma_rank) on peer's side: producer is me, slot - // is `peer.rdma_channel_head.buffer(my_rdma_rank)`). Bypasses - // both broken port_channel.put and the unreliable NVLS counter. - const int dst_rank_global = lane_id * NUM_MAX_NVL_PEERS + nvl_rank; - const uintptr_t my_head_off = reinterpret_cast(rdma_channel_head.buffer(rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - uint64_t* peer_head = - reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + my_head_off); - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(peer_head), "l"((uint64_t)min_head) : "memory"); - } else if (lane_id == rdma_rank) { - // Self-loop: plain release atomic on local slot. Cannot use NVLS - // multimem here \u2014 it fans out to all NVL peers' local buffers - // and over-counts (4 NVL ranks \u00d7 add \u21d2 4x increment). - mscclpp::atomicFetchAdd(static_cast(rdma_channel_head.buffer(rdma_rank)), - (uint64_t)(min_head - last_head), mscclpp::memoryOrderRelease); - } else { - auto dst_offset = rdma_rank * sizeof(uint64_t) + head_send_offset; - auto port_channel_idx = kLowLatencyMode ? (channel_id * kNumRDMARanks + lane_id) - : (channel_id * num_ranks + lane_id * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - // Absolute-value RDMA WRITE replaces broken HW atomicAdd (see note above). - *rdma_channel_head_send_src.buffer(lane_id) = (uint64_t)min_head; - __threadfence_system(); - const auto src_off_head = reinterpret_cast(rdma_channel_head_send_src.buffer(lane_id)) - - reinterpret_cast(rdma_buffer_ptr_base); - handle.put(dst_offset, src_off_head, sizeof(uint64_t)); - } - last_head = min_head; - } - - // Nanosleep and let other warps work - __nanosleep(NUM_WAIT_NANOSECONDS); - } - } else { - // NVL consumers - // Retrieve rank offset from barrier results (each lane's register stores an RDMA rank) - int src_nvl_rank = target_rank, total_offset = 0; - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - if (lane_id < kNumRDMARanks and lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank > 0) - total_offset = recv_gbl_rank_prefix_sum[lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank - 1]; - - // Receive channel offsets - int start_offset = 0, end_offset = 0, num_tokens_to_recv; - auto start_time = clock64(); - while (lane_id < kNumRDMARanks) { - start_offset = ld_volatile_global(nvl_channel_prefix_start.buffer() + lane_id); - end_offset = ld_volatile_global(nvl_channel_prefix_end.buffer() + lane_id); - if (start_offset < 0 and end_offset < 0) { - start_offset = -start_offset - 1, end_offset = -end_offset - 1; - total_offset += start_offset; - break; - } - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src RDMA: %d, src nvl: %d, start: " - "%d, end: %d\n", - channel_id, rdma_rank, nvl_rank, lane_id, src_nvl_rank, start_offset, end_offset); - trap(); - } - } - num_tokens_to_recv = warp_reduce_sum(end_offset - start_offset); - - // Save for combine usage - if (lane_id < kNumRDMARanks and not kCachedMode) - recv_gbl_channel_prefix_matrix[(lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank) * num_channels + channel_id] = - total_offset; - __syncwarp(); - - int cached_channel_head_idx = 0, cached_channel_tail_idx = 0; - while (num_tokens_to_recv > 0) { - // Check channel status by lane 0 - start_time = clock64(); - while (lane_id == 0) { - // Ready to copy - if (cached_channel_head_idx != cached_channel_tail_idx) break; - cached_channel_tail_idx = ld_acquire_sys_global(nvl_channel_tail.buffer()); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src NVL: %d, head: %d, tail: %d\n", - channel_id, rdma_rank, nvl_rank, src_nvl_rank, cached_channel_head_idx, cached_channel_tail_idx); - trap(); - } - } - - // Sync queue tail - cached_channel_tail_idx = __shfl_sync(0xffffffff, cached_channel_tail_idx, 0); - - // Copy data - int num_recv_tokens = cached_channel_tail_idx - cached_channel_head_idx; - for (int chunk_idx = 0; chunk_idx < num_recv_tokens; ++chunk_idx, --num_tokens_to_recv) { - int token_idx_in_buffer = (cached_channel_head_idx++) % num_max_nvl_chunked_recv_tokens; - auto meta = ld_nc_global(nvl_channel_src_meta.buffer() + token_idx_in_buffer); - int64_t recv_token_idx = __shfl_sync(0xffffffff, total_offset, meta.src_rdma_rank); - (lane_id == meta.src_rdma_rank) ? (total_offset += 1) : 0; - - // Copy data - UNROLLED_WARP_COPY(5, lane_id, hidden_int4, recv_x + recv_token_idx * hidden_int4, - nvl_channel_x.buffer() + token_idx_in_buffer * hidden_int4, ld_nc_global, st_na_global); - - // Copy source meta - if (lane_id == 0 and not kCachedMode) st_na_global(recv_src_meta + recv_token_idx, meta); - - // Copy scales - UNROLLED_WARP_COPY(1, lane_id, num_scales, recv_x_scales + recv_token_idx * num_scales, - nvl_channel_x_scales.buffer() + token_idx_in_buffer * num_scales, ld_nc_global, - st_na_global); - - // Copy `topk_idx` and `topk_weights` - if (lane_id < num_topk) { - auto recv_idx = recv_token_idx * num_topk + lane_id; - auto buffer_idx = token_idx_in_buffer * num_topk + lane_id; - st_na_global(recv_topk_idx + recv_idx, - static_cast(ld_nc_global(nvl_channel_topk_idx.buffer() + buffer_idx))); - st_na_global(recv_topk_weights + recv_idx, ld_nc_global(nvl_channel_topk_weights.buffer() + buffer_idx)); - } - } - - // Move queue - __syncwarp(); - if (lane_id == 0) st_relaxed_sys_global(nvl_channel_head.buffer(), cached_channel_head_idx); - } - } - if (thread_id == 0 && channel_id == 0) { - (void)0; - } -} - -#ifdef EP_DISPATCH_NCCLEP -#include "internode_ncclep.cuh" // warp-specialized NCCL-EP-ported dispatch_ncclep<> -#define EP_DISPATCH_KERNEL dispatch_ncclep -#define EP_DISPATCH_EXTRA_ARGS , recv_pool_ptrs, recv_pool_global_ptrs, ep_combine_recv_idx -#else -#define EP_DISPATCH_KERNEL dispatch -#define EP_DISPATCH_EXTRA_ARGS -#endif - -void dispatch(void* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv_topk_weights, void* recv_src_meta, - const void* x, const float* x_scales, const int64_t* topk_idx, const float* topk_weights, - int* send_rdma_head, int* send_nvl_head, int* recv_rdma_channel_prefix_matrix, - int* recv_gbl_channel_prefix_matrix, const int* rdma_channel_prefix_matrix, - const int* recv_rdma_rank_prefix_sum, const int* gbl_channel_prefix_matrix, - const int* recv_gbl_rank_prefix_sum, int num_tokens, int hidden_int4, int num_scales, int num_topk, - int num_experts, const bool* is_token_in_rank, void* rdma_buffer_ptr, - int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, - int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank, int num_ranks, - bool is_cached_dispatch, cudaStream_t stream, int num_channels, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_head_mc, void* nvls_head_dev, - void* nvls_tail_mc, void* nvls_tail_dev, void* const* peer_rdma_bases, void* const* recv_pool_ptrs, - void* const* recv_pool_global_ptrs, int* ep_combine_recv_idx) { - constexpr int kNumDispatchRDMASenderWarps = 6; - -#ifdef EP_DISPATCH_NCCLEP - // Increment 5 (inc5): upload MSCCLPP_EP_DIRECT -> __constant__ kEpDirect once. - { - static bool s_done_direct = false; - if (!s_done_direct) { - const char* e = std::getenv("MSCCLPP_EP_DIRECT"); - int v = (e && std::atoi(e) != 0) ? 1 : 0; - cudaMemcpyToSymbol(kEpDirect, &v, sizeof(int)); - s_done_direct = true; - } - } - // Increment 6 (inc6): upload MSCCLPP_EP_FLAT -> __constant__ kEpFlat once. - // Flat all-sender path requires kEpDirect; treated as off if direct is off. - { - static bool s_done_flat = false; - if (!s_done_flat) { - const char* ed = std::getenv("MSCCLPP_EP_DIRECT"); - const char* ef = std::getenv("MSCCLPP_EP_FLAT"); - int v = (ef && std::atoi(ef) != 0 && ed && std::atoi(ed) != 0) ? 1 : 0; - cudaMemcpyToSymbol(kEpFlat, &v, sizeof(int)); - s_done_flat = true; - } - } -#endif - -#define DISPATCH_LAUNCH_CASE(num_rdma_ranks) \ - { \ - auto dispatch_func = \ - low_latency_mode \ - ? (is_cached_dispatch ? EP_DISPATCH_KERNEL \ - : EP_DISPATCH_KERNEL) \ - : (is_cached_dispatch ? EP_DISPATCH_KERNEL \ - : EP_DISPATCH_KERNEL); \ - if (EP_NCCLEP_TMA) { \ - /* B-depth-3: opt in to the >48KB dynamic-shared half-token TMA send ring (EP_TMA_SND_*). */ \ - const size_t ep_dyn_bytes = (size_t)kNumDispatchRDMASenderWarps * EP_TMA_SND_NSTAGE * EP_TMA_SND_CHUNK_BYTES; \ - static bool s_ep_dyn_attr = false; \ - if (!s_ep_dyn_attr) { \ - CUDA_CHECK(cudaFuncSetAttribute(dispatch_func, cudaFuncAttributeMaxDynamicSharedMemorySize, \ - static_cast(ep_dyn_bytes))); \ - s_ep_dyn_attr = true; \ - } \ - cfg.dynamicSmemBytes = ep_dyn_bytes; \ - } \ - LAUNCH_KERNEL(&cfg, dispatch_func, reinterpret_cast(recv_x), recv_x_scales, recv_topk_idx, \ - recv_topk_weights, reinterpret_cast(recv_src_meta), reinterpret_cast(x), \ - x_scales, topk_idx, topk_weights, send_rdma_head, send_nvl_head, recv_rdma_channel_prefix_matrix, \ - recv_gbl_channel_prefix_matrix, rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, \ - gbl_channel_prefix_matrix, recv_gbl_rank_prefix_sum, num_tokens, hidden_int4, num_scales, num_topk, \ - num_experts, is_token_in_rank, rdma_buffer_ptr, num_max_rdma_chunked_send_tokens, \ - num_max_rdma_chunked_recv_tokens, buffer_ptrs, num_max_nvl_chunked_send_tokens, \ - num_max_nvl_chunked_recv_tokens, rank, num_ranks, port_channel_handles, memory_channel_handles, \ - nvls_head_mc, nvls_head_dev, nvls_tail_mc, nvls_tail_dev, peer_rdma_bases EP_DISPATCH_EXTRA_ARGS); \ - } \ - break - - EP_HOST_ASSERT((topk_idx == nullptr) == (topk_weights == nullptr)); - EP_HOST_ASSERT((recv_topk_idx == nullptr) == (recv_topk_weights == nullptr)); - - // inc6 (kEpFlat): the FLAT all-sender path launches 1 block per channel (no - // forwarder blocks); the default inc5 2-hop path launches 2 blocks per channel. - int ep_grid_blocks = num_channels * 2; -#ifdef EP_DISPATCH_NCCLEP - { - const char* ed = std::getenv("MSCCLPP_EP_DIRECT"); - const char* ef = std::getenv("MSCCLPP_EP_FLAT"); - if (ef && std::atoi(ef) != 0 && ed && std::atoi(ed) != 0) ep_grid_blocks = num_channels; - } -#endif - SETUP_LAUNCH_CONFIG(ep_grid_blocks, (kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32, stream); - SWITCH_RDMA_RANKS(DISPATCH_LAUNCH_CASE); -#undef DISPATCH_LAUNCH_CASE -} - -template -__global__ void cached_notify(const int rdma_clean_offset, const int rdma_num_int_clean, const int nvl_clean_offset, - const int nvl_num_int_clean, int* combined_rdma_head, int num_combined_tokens, - int num_channels, const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, - int* combined_nvl_head, void* rdma_buffer_ptr, void** buffer_ptrs, int** task_fifo_ptrs, - int head, int rank, int num_ranks, bool is_cached_dispatch, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, - // Phase 4: NVLS multimem barrier path — bypasses port_channel signal/wait - // (broken on Azure CX-7 RoCE) by using multimem.red.add + ld.acquire on a - // pair of barrier slots (offsets +24 and +32). Both nullptr ⇒ fall back to IB. - void* nvls_mc_ptr, void* nvls_dev_ptr, size_t nvls_off_barrier, uint64_t nvls_epoch) { - auto sm_id = static_cast(blockIdx.x); - auto thread_id = static_cast(threadIdx.x); - auto num_threads = static_cast(blockDim.x); - auto num_warps = num_threads / 32; - auto warp_id = thread_id / 32; - auto lane_id = get_lane_id(); - - auto rdma_rank = rank / NUM_MAX_NVL_PEERS; - auto nvl_rank = rank % NUM_MAX_NVL_PEERS; - auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; - - // Using two SMs, which clean the RDMA/NVL buffer respectively - if (sm_id == 0) { - if (thread_id == 0) { - (void)0; - } - // Barrier for RDMA — Phase 4: NVLS multimem.red.add fast path replaces - // the port_channel signal/wait pair (broken on Azure CX-7). - if (nvls_mc_ptr != nullptr) { - if (thread_id == 0) { - uint64_t* mc_b3 = reinterpret_cast(static_cast(nvls_mc_ptr) + nvls_off_barrier + 24); - uint64_t* dev_b3 = reinterpret_cast(static_cast(nvls_dev_ptr) + nvls_off_barrier + 24); - MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(mc_b3); - const uint64_t expected = nvls_epoch * static_cast(num_ranks); - uint64_t v; - do { - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(dev_b3) : "memory"); - } while (v < expected); - } - __syncthreads(); - } else { - // TODO(chhwang): it should be a global barrier when kLowLatencyMode is false - const bool run_barrier = (threadIdx.x < num_rdma_ranks) && (threadIdx.x != rdma_rank); - const auto barrier_channel_idx = kLowLatencyMode ? threadIdx.x : (threadIdx.x * NUM_MAX_NVL_PEERS + nvl_rank); - if (run_barrier) { - port_channel_handles[barrier_channel_idx].signal(); - port_channel_handles[barrier_channel_idx].wait(); - } - __syncthreads(); - } - - // Clean - auto rdma_buffer_ptr_int = reinterpret_cast(rdma_buffer_ptr); -#pragma unroll - for (int i = thread_id; i < rdma_num_int_clean; i += num_threads) rdma_buffer_ptr_int[rdma_clean_offset + i] = 0; - // Make the cleanup visible to the proxy + remote peers before the barrier. - // DeepEP used `nvshmem_fence()` here; we fall back to a system-scope - // threadfence because the actual remote visibility is provided by the - // subsequent port-channel barrier (signal + flush + wait). - __threadfence_system(); - __syncthreads(); - - // Barrier again — Phase 4: second NVLS multimem.red barrier on slot +32. - if (nvls_mc_ptr != nullptr) { - if (thread_id == 0) { - uint64_t* mc_b4 = reinterpret_cast(static_cast(nvls_mc_ptr) + nvls_off_barrier + 32); - uint64_t* dev_b4 = reinterpret_cast(static_cast(nvls_dev_ptr) + nvls_off_barrier + 32); - MSCCLPP_EP_MULTIMEM_RED_ADD_RELEASE_U64_ONE(mc_b4); - const uint64_t expected = nvls_epoch * static_cast(num_ranks); - uint64_t v; - do { - asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(v) : "l"(dev_b4) : "memory"); - } while (v < expected); - } - __syncthreads(); - } else { - const bool run_barrier = (threadIdx.x < num_rdma_ranks) && (threadIdx.x != rdma_rank); - const auto barrier_channel_idx = kLowLatencyMode ? threadIdx.x : (threadIdx.x * NUM_MAX_NVL_PEERS + nvl_rank); - if (run_barrier) { - port_channel_handles[barrier_channel_idx].signal(); - port_channel_handles[barrier_channel_idx].flush(); - port_channel_handles[barrier_channel_idx].wait(); - } - } - } else if (sm_id == 1) { - // Barrier for NVL - barrier_device(task_fifo_ptrs, head, nvl_rank); - move_fifo_slots(head); - __syncthreads(); - - // Clean - auto nvl_buffer_ptr_int = reinterpret_cast(buffer_ptrs[nvl_rank]); -#pragma unroll - for (int i = thread_id; i < nvl_num_int_clean; i += num_threads) nvl_buffer_ptr_int[nvl_clean_offset + i] = 0; - memory_fence(); - __syncthreads(); - - // Barrier again - barrier_device(task_fifo_ptrs, head, nvl_rank); - move_fifo_slots(head); - } else if (sm_id == 2) { - if (is_cached_dispatch) return; - - EP_DEVICE_ASSERT(num_rdma_ranks <= 32); - - // Iterate in reverse order. Stride over channels in warp granularity so we - // support num_channels > num_warps (per-block thread cap on GB200 is 1024). - if (lane_id < num_rdma_ranks) { - for (int ch = warp_id; ch < num_channels; ch += num_warps) { - int token_start_idx, token_end_idx; - get_channel_task_range(num_combined_tokens, num_channels, ch, token_start_idx, token_end_idx); - - // NOTES: `1 << 25` is a heuristic large number - int last_head = 1 << 25; - for (int token_idx = token_end_idx - 1; token_idx >= token_start_idx; --token_idx) { - auto current_head = __ldg(combined_rdma_head + token_idx * num_rdma_ranks + lane_id); - if (current_head < 0) { - combined_rdma_head[token_idx * num_rdma_ranks + lane_id] = -last_head - 1; - } else { - last_head = current_head; - } - } - } - } - } else { - if (is_cached_dispatch) return; - - EP_DEVICE_ASSERT(rdma_channel_prefix_matrix != nullptr and rdma_rank_prefix_sum != nullptr); - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Too many NVL peers"); - - // Stride over channels per warp so num_channels > num_warps works. - if (lane_id < NUM_MAX_NVL_PEERS) { - for (int dst_rdma_rank = sm_id - 3; dst_rdma_rank < num_rdma_ranks; dst_rdma_rank += num_channels * 2 - 3) { - for (int ch = warp_id; ch < num_channels; ch += num_warps) { - // Iterate in reverse order - int token_start_idx = ch == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + ch - 1]; - int token_end_idx = rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + ch]; - int shift = dst_rdma_rank == 0 ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1]; - token_start_idx += shift, token_end_idx += shift; - - // NOTES: `1 << 25` is a heuristic large number - int last_head = 1 << 25; - for (int token_idx = token_end_idx - 1; token_idx >= token_start_idx; --token_idx) { - auto current_head = __ldg(combined_nvl_head + token_idx * NUM_MAX_NVL_PEERS + lane_id); - if (current_head < 0) { - combined_nvl_head[token_idx * NUM_MAX_NVL_PEERS + lane_id] = -last_head - 1; - } else { - last_head = current_head; - } - } - } - } - } - } -} - -void cached_notify(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights, int num_ranks, - int num_channels, int num_combined_tokens, int* combined_rdma_head, - const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, int* combined_nvl_head, - void* rdma_buffer_ptr, int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, - int num_max_nvl_chunked_recv_tokens, int** task_fifo_ptrs, int head, int rank, cudaStream_t stream, - int64_t num_rdma_bytes, int64_t num_nvl_bytes, bool is_cached_dispatch, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_mc_ptr, void* nvls_dev_ptr, - size_t nvls_off_barrier, uint64_t nvls_epoch) { - const int num_threads = std::max(128, 32 * num_channels); - const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; - - // Get clean meta. inc5: only the DISPATCH ring is shrunk under - // MSCCLPP_EP_DIRECT (hidden goes direct to the pool); the COMBINE ring still - // carries hidden through the rdma channel, so its clean must stay full-slot. - // is_cached_dispatch distinguishes the two callers (true=cached dispatch, - // false=combine). - const bool ep_direct = []() { - const char* e = std::getenv("MSCCLPP_EP_DIRECT"); - return e && std::atoi(e) != 0; - }(); - const int clean_hidden_int4 = (ep_direct && is_cached_dispatch) ? 0 : hidden_int4; - auto rdma_clean_meta = get_rdma_clean_meta(clean_hidden_int4, num_scales, num_topk_idx, num_topk_weights, - num_rdma_ranks, num_max_rdma_chunked_recv_tokens, num_channels); - auto nvl_clean_meta = get_nvl_clean_meta(hidden_int4, num_scales, num_topk_idx, num_topk_weights, num_rdma_ranks, - NUM_MAX_NVL_PEERS, num_max_nvl_chunked_recv_tokens, num_channels); - EP_HOST_ASSERT((rdma_clean_meta.first + rdma_clean_meta.second) * sizeof(int) <= num_rdma_bytes); - EP_HOST_ASSERT((nvl_clean_meta.first + nvl_clean_meta.second) * sizeof(int) <= num_nvl_bytes); - EP_HOST_ASSERT(num_rdma_bytes < std::numeric_limits::max()); - EP_HOST_ASSERT(num_nvl_bytes < std::numeric_limits::max()); - EP_HOST_ASSERT(num_channels * 2 > 3); - - // Launch kernel. - // NOTE: cached_notify's sm_id>=2 work iterates `warp_id < num_channels`, so the - // kernel originally requested 32*num_channels threads per block. On GB200 the - // per-block thread limit is 1024 (= 32 warps), so nc>32 silently failed launch - // (cudaLaunchKernelEx returns cudaErrorInvalidValue / "invalid argument"). - // Cap at 1024 here and add a strided fallback in the kernel (warps loop over - // channels in stride-num_warps) so nc up to the SM-count is supported. - auto cached_notify_func = low_latency_mode ? cached_notify : cached_notify; - const int launch_threads = std::min(num_threads, 1024); - SETUP_LAUNCH_CONFIG(num_channels * 2, launch_threads, stream); - LAUNCH_KERNEL(&cfg, cached_notify_func, rdma_clean_meta.first, rdma_clean_meta.second, nvl_clean_meta.first, - nvl_clean_meta.second, combined_rdma_head, num_combined_tokens, num_channels, - rdma_channel_prefix_matrix, rdma_rank_prefix_sum, combined_nvl_head, rdma_buffer_ptr, buffer_ptrs, - task_fifo_ptrs, head, rank, num_ranks, is_cached_dispatch, port_channel_handles, memory_channel_handles, - nvls_mc_ptr, nvls_dev_ptr, nvls_off_barrier, nvls_epoch); -} - -template -__device__ int combine_token(bool is_token_in_rank, int head_idx, int lane_id, int hidden_int4, int num_topk, - int4* combined_row, float* combined_topk_weights, int num_max_recv_tokens, - const ReceiveFn& recv_fn, const ReceiveTWFn& recv_tw_fn) { - constexpr auto kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t); - - // Broadcast current heads - // Lane `i` holds the head of rank `i` and `is_token_in_rank` - EP_STATIC_ASSERT(kMaxNumRanks <= 32, "Too many ranks"); - int num_topk_ranks = 0, topk_ranks[kMaxNumRanks], slot_indices[kMaxNumRanks]; -#pragma unroll - for (int i = 0; i < kNumRanks; ++i) - if (__shfl_sync(0xffffffff, is_token_in_rank, i)) { - slot_indices[num_topk_ranks] = __shfl_sync(0xffffffff, head_idx, i) % num_max_recv_tokens; - topk_ranks[num_topk_ranks++] = i; - } - EP_DEVICE_ASSERT(num_topk_ranks <= kMaxNumRanks); - -// Reduce data -#pragma unroll - for (int i = lane_id; i < hidden_int4; i += 32) { - // Read buffers - // TODO: maybe too many registers here - int4 recv_value_int4[kMaxNumRanks]; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) recv_value_int4[j] = recv_fn(topk_ranks[j], slot_indices[j], i); - - // Reduce all-to-all results - float values[kDtypePerInt4] = {0}; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) { - auto recv_value_dtypes = reinterpret_cast(&recv_value_int4[j]); -#pragma unroll - for (int k = 0; k < kDtypePerInt4; ++k) values[k] += static_cast(recv_value_dtypes[k]); - } - - // Cast back to `dtype_t` and write - int4 out_int4; - auto out_dtypes = reinterpret_cast(&out_int4); -#pragma unroll - for (int j = 0; j < kDtypePerInt4; ++j) out_dtypes[j] = static_cast(values[j]); - st_na_global(combined_row + i, out_int4); - } - - // Reduce `topk_weights` - if (lane_id < num_topk) { - float value = 0; -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) value += recv_tw_fn(topk_ranks[i], slot_indices[i], lane_id); - st_na_global(combined_topk_weights + lane_id, value); - } - - // Return the minimum top-k rank - return topk_ranks[0]; -} - -// inc8 (combine extract): the FLAT direct-gather combine path, extracted from the -// unified `combine` kernel into its own lean kernel. RATIONALE: the unified kernel -// carries the heavy 2-hop forwarder code, so it compiles to ~80 registers and is -// register-limited to 1 resident block/SM (25 warps). Under flat, only this gather -// runs (the forwarder is dead code) yet it inherits that 80-reg / 1-block ceiling, -// starving the NVLink pull-gather of the warp-level parallelism it needs to hide -// the (latency-, not bandwidth-bound) remote ld_nc_global reads. Extracting it lets -// the gather compile to far fewer registers, so a full 1024-thread block (32 warps) -// fits per SM at the same SM budget -> more warps in flight -> better latency hiding. -// The token-strided loop is grid/block-agnostic, so correctness is independent of -// the launch geometry. Logic is byte-for-byte the same reduction as the unified -// kernel's flat branch (combine_token for <=8 nodes, chunked-ballot for >8). -template -__global__ void __launch_bounds__(1024, 1) - combine_flat_gather(int4* combined_x, float* combined_topk_weights, const bool* is_combined_token_in_rank, - int num_combined_tokens, int hidden, int num_topk, int num_ranks, - void* const* recv_pool_global_ptrs, const int* ep_combine_recv_idx) { - const auto lane_id = get_lane_id(); - const auto hidden_int4 = hidden / static_cast(sizeof(int4) / sizeof(dtype_t)); - constexpr int kEpNumRanks = kNumRDMARanks * NUM_MAX_NVL_PEERS; - const int64_t ep_pool_header_bytes = ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - const int warp_id_flat = static_cast(threadIdx.x) / 32; - const int num_warps_per_block = static_cast(blockDim.x) / 32; - const int global_warp = static_cast(blockIdx.x) * num_warps_per_block + warp_id_flat; - const int total_warps = static_cast(gridDim.x) * num_warps_per_block; - auto recv_fn = [&](int r, int slot, int hi) -> int4 { - return ld_nc_global(reinterpret_cast(reinterpret_cast(recv_pool_global_ptrs[r]) + - ep_pool_header_bytes) + - static_cast(slot) * hidden_int4 + hi); - }; - auto recv_tw_fn = [&](int, int, int) -> float { return 0.0f; }; - if constexpr (kEpNumRanks <= 32) { - for (int t = global_warp; t < num_combined_tokens; t += total_warps) { - const bool is_in = - (lane_id < kEpNumRanks) and is_combined_token_in_rank[static_cast(t) * num_ranks + lane_id]; - const int recv_idx = is_in ? ep_combine_recv_idx[static_cast(t) * num_ranks + lane_id] : 0; - combine_token( - is_in, recv_idx, lane_id, hidden_int4, num_topk, combined_x + static_cast(t) * hidden_int4, - combined_topk_weights + static_cast(t) * num_topk, 1 << 30, recv_fn, recv_tw_fn); - } - } else { - constexpr int kEpMaxContrib = 8; - constexpr auto kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t); - for (int t = global_warp; t < num_combined_tokens; t += total_warps) { - int topk_ranks[kEpMaxContrib], slot_indices[kEpMaxContrib], num_topk_ranks = 0; - for (int base = 0; base < kEpNumRanks; base += 32) { - const int r = base + lane_id; - const bool is_in = (r < kEpNumRanks) and is_combined_token_in_rank[static_cast(t) * num_ranks + r]; - const int slot = is_in ? ep_combine_recv_idx[static_cast(t) * num_ranks + r] : 0; - unsigned ballot = __ballot_sync(0xffffffffu, is_in); - while (ballot != 0u) { - const int l = __ffs(static_cast(ballot)) - 1; - if (num_topk_ranks < kEpMaxContrib) { - topk_ranks[num_topk_ranks] = base + l; - slot_indices[num_topk_ranks] = __shfl_sync(0xffffffffu, slot, l); - ++num_topk_ranks; - } - ballot &= ballot - 1u; - } - } - int4* combined_row = combined_x + static_cast(t) * hidden_int4; -#pragma unroll - for (int i = lane_id; i < hidden_int4; i += 32) { - int4 recv_value_int4[kEpMaxContrib]; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) recv_value_int4[j] = recv_fn(topk_ranks[j], slot_indices[j], i); - float values[kDtypePerInt4] = {0}; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) { - auto vd = reinterpret_cast(&recv_value_int4[j]); -#pragma unroll - for (int k = 0; k < kDtypePerInt4; ++k) values[k] += static_cast(vd[k]); - } - int4 out_int4; - auto out_dtypes = reinterpret_cast(&out_int4); -#pragma unroll - for (int k = 0; k < kDtypePerInt4; ++k) out_dtypes[k] = static_cast(values[k]); - st_na_global(combined_row + i, out_int4); - } - if (lane_id < num_topk) st_na_global(combined_topk_weights + static_cast(t) * num_topk + lane_id, 0.0f); - } - } -} - -// TMA-staged flat-combine gather. Stages each contributor's hidden-chunk from the remote -// NVLink recv pools into shared memory via cp.async.bulk (TMA), then reduces from SMEM and -// stores the combined token locally. This is the default flat direct-gather combine path. -// -// RATIONALE: the synchronous variant (combine_flat_gather) hides remote-NVLink read latency -// purely with memory-level parallelism + occupancy (per-thread int4[kMaxContrib] register -// batching x warps), which caps occupancy at ~50% and regresses once the grid has many -// blocks. Staging through TMA instead lets the async copy engine track many outstanding bulk -// loads at low register cost, so latency hiding no longer competes with occupancy -- it wins -// at every channel count (no crossover) and at every node scale. Topology is single-hop -// NVLink (one MNNVL LSA domain), so this is an apples-to-apples gather, not a 2-hop reduce. -// -// SHAPE: token-per-warp; the hidden row is split into kChunkInt4-sized chunks streamed through -// a kStages-deep software pipeline (prefetch kStages-1 chunks ahead). Per chunk, lane 0 posts -// one cp.async.bulk G2S per contributor into a per-warp SMEM tile; all lanes then reduce the -// chunk from SMEM. SMEM/block = kWarps * kStages * kMaxContrib * kChunkInt4 * 16 bytes -// (independent of hidden), so it requires the >48KB dynamic-shared opt-in (cudaFuncSetAttribute -// at launch). Defaults (chunk=64 int4 = 1KB descriptors, 2 stages) tuned on GB200 sm_100 @ -// hidden=7168 bf16; the warp (token) count per block is channel-adaptive (see below); -// overridable at compile time via -D for other shapes. -#ifndef EP_CMB_TMA_CHUNK_INT4 -#define EP_CMB_TMA_CHUNK_INT4 \ - 64 // hidden chunk in int4 (1KB TMA descriptors; 896 int4 / 64 = 14 chunks @ hidden=7168 bf16) -#endif -// Warp (token) count per block is CHANNEL-ADAPTIVE. More warps add token-parallelism, which wins -// when the grid has FEW blocks (low SM → many tokens/block, latency-bound), but at high block -// counts the marginal warp costs more scheduling than it buys (measured crossover: 14 warps is -// -2..-9% at <=10 channels but +3% at 16 channels). So use the WIDE count up to a channel -// threshold and the NARROW count above it. SMEM/block (WIDE 14w * 2 stages * 8 contrib * 64 -// int4 * 16B = 229KB) fits the sm_100 dynamic-shared cap. -#ifndef EP_CMB_TMA_WARPS_WIDE -#define EP_CMB_TMA_WARPS_WIDE 14 // <= EP_CMB_TMA_WARPS_MAXCH channels (low SM): more token-parallelism -#endif -#ifndef EP_CMB_TMA_WARPS_NARROW -#define EP_CMB_TMA_WARPS_NARROW 12 // > EP_CMB_TMA_WARPS_MAXCH channels (high SM): less scheduling overhead -#endif -#ifndef EP_CMB_TMA_WARPS_MAXCH -#define EP_CMB_TMA_WARPS_MAXCH 12 // channel threshold (NSM16/8ch, NSM20/10ch use WIDE; NSM32/16ch uses NARROW) -#endif -#ifndef EP_CMB_TMA_STAGES -#define EP_CMB_TMA_STAGES 2 // pipeline depth (outstanding chunks in flight) -#endif - -template -__global__ void __launch_bounds__(kWarps * 32, 1) - combine_flat_gather_tma(int4* combined_x, float* combined_topk_weights, const bool* is_combined_token_in_rank, - int num_combined_tokens, int hidden, int num_topk, int num_ranks, - void* const* recv_pool_global_ptrs, const int* ep_combine_recv_idx) { - constexpr int kMaxContrib = 8; - constexpr int kChunkInt4 = EP_CMB_TMA_CHUNK_INT4; - constexpr int kStages = EP_CMB_TMA_STAGES; // pipeline depth (outstanding chunks in flight) - constexpr int kChunkBytes = kChunkInt4 * static_cast(sizeof(int4)); - constexpr auto kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t); - constexpr int kEpNumRanks = kNumRDMARanks * NUM_MAX_NVL_PEERS; - - const auto lane_id = get_lane_id(); - const int warp_id = static_cast(threadIdx.x) / 32; - const auto hidden_int4 = hidden / static_cast(sizeof(int4) / sizeof(dtype_t)); - const int64_t ep_pool_header_bytes = ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - - // Dynamic SMEM: [kWarps][kStages][kMaxContrib][kChunkBytes] staging tiles, then [kWarps][kStages] mbarriers. - extern __shared__ uint8_t smem_raw[]; - const size_t per_warp_stage_bytes = static_cast(kStages) * kMaxContrib * kChunkBytes; - uint8_t* my_stage = smem_raw + warp_id * per_warp_stage_bytes; - uint64_t* mbar_base = reinterpret_cast(smem_raw + static_cast(kWarps) * per_warp_stage_bytes); - uint64_t* my_mbar = mbar_base + warp_id * kStages; - auto stage_buf = [&](int s, int j) -> uint8_t* { - return my_stage + (static_cast(s) * kMaxContrib + j) * kChunkBytes; - }; - - const int global_warp = static_cast(blockIdx.x) * kWarps + warp_id; - const int total_warps = static_cast(gridDim.x) * kWarps; - const int nchunks = (hidden_int4 + kChunkInt4 - 1) / kChunkInt4; - - for (int t = global_warp; t < num_combined_tokens; t += total_warps) { - // ---- discovery: which ranks contributed to this token + their recv slots ---- - int topk_ranks[kMaxContrib], slot_indices[kMaxContrib], num_topk_ranks = 0; - for (int base = 0; base < kEpNumRanks; base += 32) { - const int r = base + lane_id; - const bool is_in = (r < kEpNumRanks) and is_combined_token_in_rank[static_cast(t) * num_ranks + r]; - const int slot = is_in ? ep_combine_recv_idx[static_cast(t) * num_ranks + r] : 0; - unsigned ballot = __ballot_sync(0xffffffffu, is_in); - while (ballot != 0u) { - const int l = __ffs(static_cast(ballot)) - 1; - if (num_topk_ranks < kMaxContrib) { - topk_ranks[num_topk_ranks] = base + l; - slot_indices[num_topk_ranks] = __shfl_sync(0xffffffffu, slot, l); - ++num_topk_ranks; - } - ballot &= ballot - 1u; - } - } - - int4* combined_row = combined_x + static_cast(t) * hidden_int4; - - // lane 0 posts all contributors' TMA G2S loads for one chunk into stage s. - auto issue_loads = [&](int s, int c0, int csize_int4) { - if (lane_id == 0) { - const uint32_t mbar_a = static_cast(__cvta_generic_to_shared(&my_mbar[s])); - asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(mbar_a)); - fenceProxyAsyncSharedCta(); - const uint32_t cbytes = static_cast(csize_int4 * static_cast(sizeof(int4))); - for (int j = 0; j < num_topk_ranks; ++j) { - const uint8_t* src = - reinterpret_cast(recv_pool_global_ptrs[topk_ranks[j]]) + ep_pool_header_bytes + - static_cast(slot_indices[j]) * hidden_int4 * static_cast(sizeof(int4)) + - static_cast(c0) * sizeof(int4); - const uint32_t dst = static_cast(__cvta_generic_to_shared(stage_buf(s, j))); - asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst), - "l"(src), "r"(cbytes), "r"(mbar_a) - : "memory"); - } - uint64_t st; - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 %0, [%1], %2;" - : "=l"(st) - : "r"(mbar_a), "r"(static_cast(cbytes * num_topk_ranks))); - } - }; - - // all lanes wait for stage s's loads, then fence so generic reads see the async-proxy writes. - auto wait_stage = [&](int s) { - if (lane_id == 0) { - const uint32_t mbar_a = static_cast(__cvta_generic_to_shared(&my_mbar[s])); - uint32_t done = 0; - while (!done) { - asm volatile( - "{ .reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], 0;" - " selp.u32 %0, 1, 0, p; }" - : "=r"(done) - : "r"(mbar_a)); - } - } - __syncwarp(); - fenceProxyAsyncSharedCta(); - }; - - auto reduce_store = [&](int s, int c0, int csize_int4) { - for (int i = lane_id; i < csize_int4; i += 32) { - float values[kDtypePerInt4]; -#pragma unroll - for (int k = 0; k < static_cast(kDtypePerInt4); ++k) values[k] = 0.0f; -#pragma unroll - for (int j = 0; j < kMaxContrib; ++j) { - if (j >= num_topk_ranks) break; - const int4 v = *reinterpret_cast(stage_buf(s, j) + i * static_cast(sizeof(int4))); - auto vd = reinterpret_cast(&v); -#pragma unroll - for (int k = 0; k < static_cast(kDtypePerInt4); ++k) values[k] += static_cast(vd[k]); - } - int4 out_int4; - auto out_d = reinterpret_cast(&out_int4); -#pragma unroll - for (int k = 0; k < static_cast(kDtypePerInt4); ++k) out_d[k] = static_cast(values[k]); - st_na_global(combined_row + c0 + i, out_int4); - } - }; - - // Software pipeline: prologue issues the first kStages-1 chunks; each iteration issues - // chunk c+(kStages-1) while waiting on + reducing chunk c. Each stage has its own SMEM - // buffer + mbarrier, so up to kStages-1 chunks (x num_topk_ranks TMA loads each) stay in - // flight, giving the async copy engine the memory-level parallelism cmbext gets from its - // per-thread int4[8] batching -- but at low register cost. -#pragma unroll - for (int p = 0; p < kStages - 1; ++p) { - if (p < nchunks) { - const int c0 = p * kChunkInt4; - const int csize = (c0 + kChunkInt4 <= hidden_int4) ? kChunkInt4 : (hidden_int4 - c0); - issue_loads(p % kStages, c0, csize); - } - } - for (int c = 0; c < nchunks; ++c) { - const int s = c % kStages; - const int c0 = c * kChunkInt4; - const int csize = (c0 + kChunkInt4 <= hidden_int4) ? kChunkInt4 : (hidden_int4 - c0); - const int cn = c + (kStages - 1); - if (cn < nchunks) { - const int sn = cn % kStages; - const int c0n = cn * kChunkInt4; - const int csizen = (c0n + kChunkInt4 <= hidden_int4) ? kChunkInt4 : (hidden_int4 - c0n); - issue_loads(sn, c0n, csizen); - } - wait_stage(s); - reduce_store(s, c0, csize); - __syncwarp(); - } - if (lane_id < num_topk) st_na_global(combined_topk_weights + static_cast(t) * num_topk + lane_id, 0.0f); - } -} - -template 0) ? kNumCombineForwarderWarps / kNumRDMARanks : 1, - int kNumForwarders = kNumRDMARanks * kNumWarpsPerForwarder, - int kNumRDMAReceivers = kNumForwarders + NUM_MAX_NVL_PEERS> -__global__ void __launch_bounds__((NUM_MAX_NVL_PEERS + 1 + kNumForwarders) * 32, 1) - combine(int4* combined_x, float* combined_topk_weights, const bool* is_combined_token_in_rank, const int4* x, - const float* topk_weights, const int* combined_rdma_head, const int* combined_nvl_head, - const SourceMeta* src_meta, const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, - const int* gbl_channel_prefix_matrix, int num_tokens, int num_combined_tokens, int hidden, int num_topk, - void* rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens, - void** buffer_ptrs, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank, - int num_ranks, mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, - // Phase 3 NVLS counter pointers (nullptr → fall back to PortChannel/atomicAdd path). - void* nvls_head_mc, void* nvls_head_dev, void* nvls_tail_mc, void* nvls_tail_dev, - // Phase 4 fabric-IPC peer base pointers (nullptr → fall back to handle.put for cross-node data). - void* const* peer_rdma_bases, - // Increment 5 combine-direct: domain-wide recv-pool bases (by global rank) + - // the dispatch gather map. Non-null + kEpDirect => combine gathers each token's - // contributions straight from the peer pools and reduces locally (no 2-hop). - void* const* recv_pool_global_ptrs, const int* ep_combine_recv_idx) { - enum class WarpRole { kNVLSender, kNVLAndRDMAForwarder, kRDMAReceiver, kCoordinator }; - - const auto sm_id = static_cast(blockIdx.x); - const auto thread_id = static_cast(threadIdx.x), lane_id = get_lane_id(); - const auto num_channels = static_cast(gridDim.x) / 2, channel_id = sm_id / 2; - const bool is_rdma_receiver_sm = sm_id % 2 == 1; - - EP_DEVICE_ASSERT(num_topk <= 32); - EP_DEVICE_ASSERT(hidden % (sizeof(int4) / sizeof(dtype_t)) == 0); - const auto hidden_int4 = hidden / (sizeof(int4) / sizeof(dtype_t)); - - // NOTES: we decouple a channel into 2 SMs - const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; - if (rank == 0 && thread_id == 0 && channel_id == 0) { - (void)0; - } -#ifdef EP_DISPATCH_NCCLEP - // inc5 combine-direct: under MSCCLPP_EP_DIRECT, gather each combined token's - // contributions DIRECTLY from the peer recv pools (recv_pool_global_ptrs[r] + - // header + recv_idx*hidden) and reduce locally, skipping nvl_channel + - // forwarder + rdma_channel. recv_idx is the dispatch gather map - // (ep_combine_recv_idx[t*num_ranks+r]); this mirrors dispatch's sender-direct - // write in reverse. Assumes the combine input lives in the recv pool (DeepEP - // contract; the flat test passes recv_x). - if (kEpDirect and recv_pool_global_ptrs != nullptr and ep_combine_recv_idx != nullptr) { - // kEpNumRanks is the world size for this kernel instantiation (compile-time: - // kNumRDMARanks * NUM_MAX_NVL_PEERS, == num_ranks). The fast combine_token<> - // discovery maps one warp lane per rank; that is correct AND fast only while - // kEpNumRanks <= 32. For >= 16 nodes (kEpNumRanks > 32) its __shfl-by-rank source - // lane wraps mod 32 (ranks >= 32 alias 0..31 -> double-count / garbage), so we fall - // back to a chunked-ballot discovery that scans ranks in 32-wide chunks. The branch - // is `if constexpr`, so <= 8 nodes compile to exactly the original fast path. - constexpr int kEpNumRanks = kNumRDMARanks * NUM_MAX_NVL_PEERS; - const int64_t ep_pool_header_bytes = ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - const int warp_id_flat = thread_id / 32; - const int num_warps_per_block = static_cast(blockDim.x) / 32; - const int global_warp = sm_id * num_warps_per_block + warp_id_flat; - const int total_warps = static_cast(gridDim.x) * num_warps_per_block; - auto recv_fn = [&](int r, int slot, int hi) -> int4 { - return ld_nc_global(reinterpret_cast(reinterpret_cast(recv_pool_global_ptrs[r]) + - ep_pool_header_bytes) + - static_cast(slot) * hidden_int4 + hi); - }; - auto recv_tw_fn = [&](int, int, int) -> float { return 0.0f; }; - if constexpr (kEpNumRanks <= 32) { - // Original fast path: one lane per rank; combine_token does compile-time-unrolled - // discovery + register-array reduction. Correct because kEpNumRanks <= 32. - for (int t = global_warp; t < num_combined_tokens; t += total_warps) { - const bool is_in = - (lane_id < kEpNumRanks) and is_combined_token_in_rank[static_cast(t) * num_ranks + lane_id]; - const int recv_idx = is_in ? ep_combine_recv_idx[static_cast(t) * num_ranks + lane_id] : 0; - combine_token( - is_in, recv_idx, lane_id, hidden_int4, num_topk, combined_x + static_cast(t) * hidden_int4, - combined_topk_weights + static_cast(t) * num_topk, 1 << 30, recv_fn, recv_tw_fn); - } - } else { - // >= 16 nodes (kEpNumRanks > 32): chunked-ballot discovery (lane-per-rank aliases - // past 32). A token routes to <= num_topk (<= 8) distinct ranks => <= 8 contributors. - constexpr int kEpMaxContrib = 8; - constexpr auto kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t); - for (int t = global_warp; t < num_combined_tokens; t += total_warps) { - // Scan ranks in chunks of 32 and compact the set lanes via ballot. Discovery is - // warp-uniform: every lane ends with the same topk_ranks/slot_indices/num_topk_ranks. - int topk_ranks[kEpMaxContrib], slot_indices[kEpMaxContrib], num_topk_ranks = 0; - for (int base = 0; base < kEpNumRanks; base += 32) { - const int r = base + lane_id; - const bool is_in = (r < kEpNumRanks) and is_combined_token_in_rank[static_cast(t) * num_ranks + r]; - const int slot = is_in ? ep_combine_recv_idx[static_cast(t) * num_ranks + r] : 0; - unsigned ballot = __ballot_sync(0xffffffffu, is_in); - while (ballot != 0u) { - const int l = __ffs(static_cast(ballot)) - 1; - if (num_topk_ranks < kEpMaxContrib) { - topk_ranks[num_topk_ranks] = base + l; - slot_indices[num_topk_ranks] = __shfl_sync(0xffffffffu, slot, l); - ++num_topk_ranks; - } - ballot &= ballot - 1u; - } - } - // Reduce the contributors' hidden rows into combined_x (lane-strided), mirroring - // combine_token: #pragma unroll the outer hidden loop + pre-load contributors. - int4* combined_row = combined_x + static_cast(t) * hidden_int4; -#pragma unroll - for (int i = lane_id; i < hidden_int4; i += 32) { - int4 recv_value_int4[kEpMaxContrib]; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) recv_value_int4[j] = recv_fn(topk_ranks[j], slot_indices[j], i); - float values[kDtypePerInt4] = {0}; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) { - auto vd = reinterpret_cast(&recv_value_int4[j]); -#pragma unroll - for (int k = 0; k < kDtypePerInt4; ++k) values[k] += static_cast(vd[k]); - } - int4 out_int4; - auto out_dtypes = reinterpret_cast(&out_int4); -#pragma unroll - for (int k = 0; k < kDtypePerInt4; ++k) out_dtypes[k] = static_cast(values[k]); - st_na_global(combined_row + i, out_int4); - } - // combined_topk_weights: gather path contributes 0 (parity with recv_tw_fn). - if (lane_id < num_topk) - st_na_global(combined_topk_weights + static_cast(t) * num_topk + lane_id, 0.0f); - } - } - return; - } -#endif - auto role_meta = [=]() -> std::pair { - auto warp_id = thread_id / 32; - if (not is_rdma_receiver_sm) { - if (warp_id < NUM_MAX_NVL_PEERS) { - auto shuffled_warp_id = warp_id; - shuffled_warp_id = (shuffled_warp_id + channel_id) % NUM_MAX_NVL_PEERS; - return {WarpRole::kNVLSender, shuffled_warp_id}; - } else if (warp_id < NUM_MAX_NVL_PEERS + kNumForwarders) { - auto shuffled_warp_id = warp_id - NUM_MAX_NVL_PEERS; - shuffled_warp_id = (shuffled_warp_id + channel_id) % kNumForwarders; - return {WarpRole::kNVLAndRDMAForwarder, shuffled_warp_id}; - } else { - return {WarpRole::kCoordinator, 0}; - } - } else { - if (warp_id < NUM_MAX_NVL_PEERS + kNumForwarders) { - return {WarpRole::kRDMAReceiver, warp_id}; - } else { - return {WarpRole::kCoordinator, 0}; - } - } - }(); - auto warp_role = role_meta.first; - auto warp_id = role_meta.second; - - EP_DEVICE_ASSERT(num_warps == NUM_MAX_NVL_PEERS + kNumForwarders + 1); - auto num_max_nvl_chunked_recv_tokens_per_rdma = num_max_nvl_chunked_recv_tokens / kNumRDMARanks; - if (thread_id == 0 && channel_id == 0 && rank == 0) { - (void)0; - } - - if (warp_role == WarpRole::kNVLSender) { - // NVL producers - const auto dst_nvl_rank = warp_id; - - // NVL layouts - // NOTES: to avoid deadlocks, we use separate NVL buffers for different RDMA sources - auto dst_buffer_ptr = buffer_ptrs[dst_nvl_rank], local_buffer_ptr = buffer_ptrs[nvl_rank]; - auto nvl_channel_x = AsymBuffer(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens * hidden_int4, - NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) - .advance_also(local_buffer_ptr); - auto nvl_channel_src_meta = AsymBuffer(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens, - NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) - .advance_also(local_buffer_ptr); - auto nvl_channel_topk_weights = AsymBuffer(dst_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk, - NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) - .advance_also(local_buffer_ptr); - auto nvl_channel_head = - AsymBuffer(local_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, dst_nvl_rank) - .advance_also(dst_buffer_ptr); - auto nvl_channel_tail = - AsymBuffer(dst_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, nvl_rank) - .advance_also(local_buffer_ptr); - - // Get tasks for each RDMA lane - int token_start_idx = 0, token_end_idx = 0; - if (lane_id < kNumRDMARanks) { - int prefix_idx = (lane_id * NUM_MAX_NVL_PEERS + dst_nvl_rank) * num_channels + channel_id; - token_start_idx = gbl_channel_prefix_matrix[prefix_idx]; - token_end_idx = - (prefix_idx == num_channels * num_ranks - 1) ? num_tokens : gbl_channel_prefix_matrix[prefix_idx + 1]; - } - __syncwarp(); - - // NOTES: here the cached value of each lane is only responsible for a single RDMA buffer - int cached_channel_head_idx = 0, cached_channel_tail_idx = 0; - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - - // Iterate over all tokens and send by chunks - while (true) { - // Exit if possible - if (__all_sync(0xffffffff, token_start_idx >= token_end_idx)) break; - - // Decide next RDMA buffer to send - bool is_lane_ready = false; - auto start_time = clock64(); - while (true) { - int num_used_slots = cached_channel_tail_idx - cached_channel_head_idx; - is_lane_ready = lane_id < kNumRDMARanks and token_start_idx < token_end_idx and - num_max_nvl_chunked_recv_tokens_per_rdma - num_used_slots >= num_max_nvl_chunked_send_tokens; - if (__any_sync(0xffffffff, is_lane_ready)) break; - - // Retry - if (lane_id < kNumRDMARanks and token_start_idx < token_end_idx) - cached_channel_head_idx = ld_volatile_global(nvl_channel_head.buffer() + lane_id); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { - printf( - "DeepEP combine NVL sender timeout, channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, RDMA lane: %d, head: " - "%d, tail: %d, start: %d, end: %d\n", - channel_id, rdma_rank, nvl_rank, dst_nvl_rank, lane_id, - ld_volatile_global(nvl_channel_head.buffer() + lane_id), cached_channel_tail_idx, token_start_idx, - token_end_idx); - trap(); - } - } - - // Sync token start index and count - for (int current_rdma_idx = 0; current_rdma_idx < kNumRDMARanks; ++current_rdma_idx) { - if (__shfl_sync(0xffffffff, (token_start_idx >= token_end_idx) or (not is_lane_ready), current_rdma_idx)) - continue; - - // Sync token start index - auto token_idx = static_cast(__shfl_sync(0xffffffff, token_start_idx, current_rdma_idx)); - int num_tokens_in_chunk = __shfl_sync( - 0xffffffff, min(num_max_nvl_chunked_send_tokens, token_end_idx - token_start_idx), current_rdma_idx); - - // Send by chunk - for (int chunk_idx = 0; chunk_idx < num_tokens_in_chunk; ++chunk_idx, ++token_idx) { - // Get an empty slot - int dst_slot_idx = 0; - if (lane_id == current_rdma_idx) { - dst_slot_idx = (cached_channel_tail_idx++) % num_max_nvl_chunked_recv_tokens_per_rdma; - dst_slot_idx = current_rdma_idx * num_max_nvl_chunked_recv_tokens_per_rdma + dst_slot_idx; - } - dst_slot_idx = __shfl_sync(0xffffffff, dst_slot_idx, current_rdma_idx); - - // Copy data - auto shifted_x_buffers = nvl_channel_x.buffer() + dst_slot_idx * hidden_int4; - auto shifted_x = x + token_idx * hidden_int4; - UNROLLED_WARP_COPY(5, lane_id, hidden_int4, shifted_x_buffers, shifted_x, ld_nc_global, st_na_global); - - // Copy source meta - if (lane_id == 0) - st_na_global(nvl_channel_src_meta.buffer() + dst_slot_idx, ld_nc_global(src_meta + token_idx)); - - // Copy `topk_weights` - if (lane_id < num_topk) - st_na_global(nvl_channel_topk_weights.buffer() + dst_slot_idx * num_topk + lane_id, - ld_nc_global(topk_weights + token_idx * num_topk + lane_id)); - } - lane_id == current_rdma_idx ? (token_start_idx = static_cast(token_idx)) : 0; - } - - // Move queue tail - __syncwarp(); - if (lane_id < kNumRDMARanks and is_lane_ready) - st_release_sys_global(nvl_channel_tail.buffer() + lane_id, cached_channel_tail_idx); - } - } else { - // Combiners and coordinators - // RDMA symmetric layout (snapshot the base before SymBuffer advances it). - void* const rdma_buffer_ptr_base = rdma_buffer_ptr; - auto hidden_bytes = hidden_int4 * sizeof(int4); - auto num_bytes_per_rdma_token = get_num_bytes_per_rdma_token(hidden_int4, 0, 0, num_topk); - auto rdma_channel_data = - SymBuffer(rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token, kNumRDMARanks, - channel_id, num_channels); - auto rdma_channel_head = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_tail = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - // Scratch slots for absolute-value RDMA WRITE replacement of broken HW - // atomicAdd on Azure CX-7 RoCE (see dispatch kernel for full rationale). - auto rdma_channel_tail_send_src = - SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_head_send_src = - SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - - auto data_send_offset = - sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * kNumRDMARanks * channel_id; - auto data_recv_offset = sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * - kNumRDMARanks * (channel_id + num_channels); - auto head_offset = sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * kNumRDMARanks * - num_channels * 2; - auto head_send_offset = head_offset + sizeof(uint64_t) * kNumRDMARanks * channel_id; - auto tail_offset = head_offset + sizeof(uint64_t) * kNumRDMARanks * num_channels; - auto tail_send_offset = tail_offset + sizeof(uint64_t) * kNumRDMARanks * channel_id; - - // NVL layouts - void* local_nvl_buffer = buffer_ptrs[nvl_rank]; - void* nvl_buffers[NUM_MAX_NVL_PEERS]; -#pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) nvl_buffers[i] = buffer_ptrs[i]; - auto nvl_channel_x = AsymBuffer(local_nvl_buffer, num_max_nvl_chunked_recv_tokens * hidden_int4, - NUM_MAX_NVL_PEERS, channel_id, num_channels) - .advance_also(nvl_buffers); - auto nvl_channel_src_meta = AsymBuffer(local_nvl_buffer, num_max_nvl_chunked_recv_tokens, - NUM_MAX_NVL_PEERS, channel_id, num_channels) - .advance_also(nvl_buffers); - auto nvl_channel_topk_weights = AsymBuffer(local_nvl_buffer, num_max_nvl_chunked_recv_tokens * num_topk, - NUM_MAX_NVL_PEERS, channel_id, num_channels) - .advance_also(nvl_buffers); - auto nvl_channel_head = AsymBuffer(nvl_buffers, kNumRDMARanks, NUM_MAX_NVL_PEERS, - channel_id, num_channels, nvl_rank) - .advance_also(local_nvl_buffer); - auto nvl_channel_tail = - AsymBuffer(local_nvl_buffer, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels) - .advance_also(nvl_buffers); - - // Combiner warp synchronization - __shared__ volatile int forwarder_nvl_head[kNumForwarders][NUM_MAX_NVL_PEERS]; - __shared__ volatile bool forwarder_retired[kNumForwarders]; - __shared__ volatile int rdma_receiver_rdma_head[kNumRDMAReceivers][kNumRDMARanks]; - __shared__ volatile bool rdma_receiver_retired[kNumRDMAReceivers]; - auto sync_forwarder_smem = [=]() { asm volatile("bar.sync 0, %0;" ::"r"((kNumForwarders + 1) * 32)); }; - auto sync_rdma_receiver_smem = [=]() { asm volatile("bar.sync 1, %0;" ::"r"((kNumRDMAReceivers + 1) * 32)); }; - - if (warp_role == WarpRole::kNVLAndRDMAForwarder) { - // Receive from NVL ranks and forward to RDMA ranks - // NOTES: this part is using "large warps" for each RDMA ranks - const auto dst_rdma_rank = warp_id / kNumWarpsPerForwarder; - const auto sub_warp_id = warp_id % kNumWarpsPerForwarder; - auto send_buffer = dst_rdma_rank == rdma_rank ? rdma_channel_data.recv_buffer(dst_rdma_rank) - : rdma_channel_data.send_buffer(dst_rdma_rank); - auto sync_large_warp = [=]() { - if (kNumWarpsPerForwarder == 1) { - __syncwarp(); - } else { - asm volatile("bar.sync %0, %1;" ::"r"(dst_rdma_rank + 2), "r"(kNumWarpsPerForwarder * 32)); - } - }; - EP_STATIC_ASSERT(kNumWarpsPerForwarder == 1 or kNumRDMARanks + 2 <= 16, "Barriers are not enough"); - - // Advance to the corresponding NVL buffer - nvl_channel_x.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma * hidden_int4); - nvl_channel_src_meta.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma); - nvl_channel_topk_weights.advance(dst_rdma_rank * num_max_nvl_chunked_recv_tokens_per_rdma * num_topk); - nvl_channel_head.advance(dst_rdma_rank); - nvl_channel_tail.advance(dst_rdma_rank); - - // Clean shared memory and sync - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); - lane_id < NUM_MAX_NVL_PEERS ? (forwarder_nvl_head[warp_id][lane_id] = 0) : 0; - lane_id == 0 ? (forwarder_retired[warp_id] = false) : false; - sync_forwarder_smem(); - - // Get count and cached head - int cached_nvl_channel_tail_idx = 0; - int num_tokens_to_combine = rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id]; - int num_tokens_prefix = - channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1]; - num_tokens_to_combine -= num_tokens_prefix; - num_tokens_prefix += dst_rdma_rank == 0 ? 0 : rdma_rank_prefix_sum[dst_rdma_rank - 1]; - combined_nvl_head += num_tokens_prefix * NUM_MAX_NVL_PEERS; - - // Iterate over all tokens and combine by chunks - for (int token_start_idx = 0; token_start_idx < num_tokens_to_combine; - token_start_idx += num_max_rdma_chunked_send_tokens) { - // Check destination queue emptiness, or wait a buffer to be released - auto token_end_idx = min(token_start_idx + num_max_rdma_chunked_send_tokens, num_tokens_to_combine); - auto num_chunked_tokens = token_end_idx - token_start_idx; - auto start_time = clock64(); - while (sub_warp_id == 0 and lane_id == 0) { - // Inequality: `num_max_rdma_chunked_recv_tokens - (tail - head) >= num_chunked_tokens` - // Here, `token_start_idx` is the actual tail - // Phase 4: head read — cross-node head feedback comes from peer - // via fabric-VA store, self-loop head from local atomic. Both end - // up in rdma_channel_head; one read path covers them. - int num_used_slots = - token_start_idx - static_cast(ld_volatile_global(rdma_channel_head.buffer(dst_rdma_rank))); - if (num_max_rdma_chunked_recv_tokens - num_used_slots >= num_chunked_tokens) break; - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP combine forwarder (RDMA check) timeout, channel: %d, RDMA: %d, nvl: %d, dst RDMA: %d, head: " - "%ld, tail: %d, chunked: %d\n", - channel_id, rdma_rank, nvl_rank, dst_rdma_rank, - ld_volatile_global(rdma_channel_head.buffer(dst_rdma_rank)), token_start_idx, num_chunked_tokens); - trap(); - } - } - sync_large_warp(); - - // Combine and write to the RDMA buffer - for (int token_idx = token_start_idx + sub_warp_id; token_idx < token_end_idx; - token_idx += kNumWarpsPerForwarder) { - // Read expected head - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - int expected_head = -1; - if (lane_id < NUM_MAX_NVL_PEERS) - expected_head = ld_nc_global(combined_nvl_head + token_idx * NUM_MAX_NVL_PEERS + lane_id); - - // Wait lanes to be ready - start_time = clock64(); - while (cached_nvl_channel_tail_idx <= expected_head) { - cached_nvl_channel_tail_idx = ld_acquire_sys_global(nvl_channel_tail.buffer(lane_id)); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < NUM_MAX_NVL_PEERS) { - printf( - "DeepEP combine forwarder (NVL check) timeout, channel: %d, RDMA: %d, nvl: %d, src NVL: %d, dst " - "RDMA: %d, tail: %d, waiting: %d, total: %d, sub: %d, large: %d, expected: %d\n", - channel_id, rdma_rank, nvl_rank, lane_id, dst_rdma_rank, cached_nvl_channel_tail_idx, token_idx, - num_tokens_to_combine, sub_warp_id, kNumWarpsPerForwarder, expected_head); - trap(); - } - } - - // Combine current token - auto rdma_slot_idx = token_idx % num_max_rdma_chunked_recv_tokens; - void* shifted = send_buffer + rdma_slot_idx * num_bytes_per_rdma_token; - auto recv_fn = [&](int src_nvl_rank, int slot_idx, int hidden_int4_idx) -> int4 { - return ld_nc_global(nvl_channel_x.buffer(src_nvl_rank) + slot_idx * hidden_int4 + hidden_int4_idx); - }; - auto recv_tw_fn = [&](int src_nvl_rank, int slot_idx, int topk_idx) -> float { - return ld_nc_global(nvl_channel_topk_weights.buffer(src_nvl_rank) + slot_idx * num_topk + topk_idx); - }; - combine_token( - expected_head >= 0, expected_head, lane_id, hidden_int4, num_topk, reinterpret_cast(shifted), - reinterpret_cast(reinterpret_cast(shifted) + hidden_bytes + sizeof(SourceMeta)), - num_max_nvl_chunked_recv_tokens_per_rdma, recv_fn, recv_tw_fn); - - // Update head - if (lane_id < NUM_MAX_NVL_PEERS) - expected_head < 0 ? (forwarder_nvl_head[warp_id][lane_id] = -expected_head - 1) - : (forwarder_nvl_head[warp_id][lane_id] = expected_head + 1); - } - sync_large_warp(); - - // Issue RDMA send - if (sub_warp_id == kNumWarpsPerForwarder - 1) { - if (nvls_tail_mc != nullptr) { - // Phase 3+4: NVLS counter fast path. Cross-node payload uses - // warp-cooperative direct fabric-IPC writes when peer_rdma_bases - // is available; otherwise single-lane handle.put + flush. - if (dst_rdma_rank != rdma_rank) { - const auto rdma_slot_idx = token_start_idx % num_max_rdma_chunked_recv_tokens; - const size_t num_bytes_per_msg = (size_t)num_chunked_tokens * (size_t)num_bytes_per_rdma_token; - const auto dst_offset = rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - rdma_slot_idx * num_bytes_per_rdma_token + data_recv_offset; - const auto src_offset = dst_rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - rdma_slot_idx * num_bytes_per_rdma_token + data_send_offset; - if (peer_rdma_bases != nullptr) { - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const int4* src_p = - reinterpret_cast(reinterpret_cast(rdma_buffer_ptr_base) + src_offset); - int4* dst_p = - reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + dst_offset); - const int n_int4 = (int)(num_bytes_per_msg / sizeof(int4)); - // Same 8x unroll as dispatch sender for LSU pipelining. - const int stride8 = 8 * 32; - int k = lane_id; - const int n_full = (n_int4 / stride8) * stride8; - for (; k < n_full; k += stride8) { - int4 v0 = src_p[k]; - int4 v1 = src_p[k + 32]; - int4 v2 = src_p[k + 64]; - int4 v3 = src_p[k + 96]; - int4 v4 = src_p[k + 128]; - int4 v5 = src_p[k + 160]; - int4 v6 = src_p[k + 192]; - int4 v7 = src_p[k + 224]; - dst_p[k] = v0; - dst_p[k + 32] = v1; - dst_p[k + 64] = v2; - dst_p[k + 96] = v3; - dst_p[k + 128] = v4; - dst_p[k + 160] = v5; - dst_p[k + 192] = v6; - dst_p[k + 224] = v7; - } - for (; k < n_int4; k += 32) { - dst_p[k] = src_p[k]; - } - __syncwarp(); - __threadfence_system(); - } else if (lane_id == 0) { - const auto port_channel_idx = - kLowLatencyMode ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - handle.put(dst_offset, src_offset, num_bytes_per_msg); - handle.flush(); - } - } - // Phase 4: tail counter publish. - // cross-node → direct fabric-VA st.release on peer's tail slot - // self-loop → plain local atomicAdd (NVLS multicast over-counts - // across NVL peers — see dispatch fix). - if (dst_rdma_rank != rdma_rank) { - if (peer_rdma_bases != nullptr && lane_id == 0) { - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const uintptr_t my_tail_off = reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - uint64_t* peer_tail = reinterpret_cast( - reinterpret_cast(peer_rdma_bases[dst_rank_global]) + my_tail_off); - const uint64_t new_tail = (uint64_t)(token_start_idx + num_chunked_tokens); - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(peer_tail), "l"(new_tail) : "memory"); - } - } else if (lane_id == 0) { - mscclpp::atomicFetchAdd(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)), - (uint64_t)num_chunked_tokens, mscclpp::memoryOrderRelease); - } - } else if (lane_id == 0) { - if (dst_rdma_rank == rdma_rank) { - mscclpp::atomicFetchAdd(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)), - (uint64_t)num_chunked_tokens, mscclpp::memoryOrderRelease); - } else { - auto rdma_slot_idx = token_start_idx % num_max_rdma_chunked_recv_tokens; - const size_t num_bytes_per_msg = num_chunked_tokens * num_bytes_per_rdma_token; - auto dst_offset = rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - rdma_slot_idx * num_bytes_per_rdma_token + data_recv_offset; - auto src_offset = dst_rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - rdma_slot_idx * num_bytes_per_rdma_token + data_send_offset; - auto port_channel_idx = kLowLatencyMode - ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - handle.put(dst_offset, src_offset, num_bytes_per_msg); - - // Absolute-value RDMA WRITE replaces broken HW atomicAdd. - const uint64_t new_tail = (uint64_t)(token_start_idx + num_chunked_tokens); - *rdma_channel_tail_send_src.buffer(dst_rdma_rank) = new_tail; - __threadfence_system(); - const auto src_off_tail = reinterpret_cast(rdma_channel_tail_send_src.buffer(dst_rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - handle.put(rdma_rank * sizeof(uint64_t) + tail_send_offset, src_off_tail, sizeof(uint64_t)); - } - } - __syncwarp(); - } - } - - // Retired - __syncwarp(); - if (lane_id == 0) forwarder_retired[warp_id] = true; - } else if (warp_role == WarpRole::kRDMAReceiver) { - // Receive from RDMA ranks and write to the output tensor - // Clean shared memory and sync - EP_DEVICE_ASSERT(kNumRDMARanks <= 32); - lane_id < kNumRDMARanks ? (rdma_receiver_rdma_head[warp_id][lane_id] = 0) : 0; - lane_id == 0 ? (rdma_receiver_retired[warp_id] = false) : 0; - sync_rdma_receiver_smem(); - - // The same tokens as the dispatch process - int token_start_idx, token_end_idx; - get_channel_task_range(num_combined_tokens, num_channels, channel_id, token_start_idx, token_end_idx); - - // Iterate over all tokens and combine - int cached_channel_tail_idx = 0; - for (int64_t token_idx = token_start_idx + warp_id; token_idx < token_end_idx; token_idx += kNumRDMAReceivers) { - // Read expected head - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - int expected_head = -1; - if (lane_id < kNumRDMARanks) { - expected_head = ld_nc_global(combined_rdma_head + token_idx * kNumRDMARanks + lane_id); - (expected_head < 0) ? (rdma_receiver_rdma_head[warp_id][lane_id] = -expected_head - 1) - : (rdma_receiver_rdma_head[warp_id][lane_id] = expected_head); - } - - // Wait lanes to be ready - auto start_time = clock64(); - while (cached_channel_tail_idx <= expected_head) { - // Phase 4: receiver waits on rdma_channel_tail. Cross-node sender - // wrote it via fabric-VA, self sender wrote it via local atomic. - // One read path via ld_acquire_sys_global covers both. - cached_channel_tail_idx = static_cast(ld_acquire_sys_global(rdma_channel_tail.buffer(lane_id))); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP combine RDMA receiver timeout, channel: %d, RDMA: %d, nvl: %d, src RDMA: %d, tail: %d, " - "waiting: %ld, expect: %d\n", - channel_id, rdma_rank, nvl_rank, lane_id, cached_channel_tail_idx, token_idx, expected_head); - trap(); - } - } - __syncwarp(); - - // Combine current token - auto recv_fn = [&](int src_rdma_rank, int slot_idx, int hidden_int4_idx) -> int4 { - return ld_nc_global(reinterpret_cast(rdma_channel_data.recv_buffer(src_rdma_rank) + - slot_idx * num_bytes_per_rdma_token) + - hidden_int4_idx); - }; - auto recv_tw_fn = [&](int src_rdma_rank, int slot_idx, int topk_idx) -> float { - return ld_nc_global(reinterpret_cast(rdma_channel_data.recv_buffer(src_rdma_rank) + - slot_idx * num_bytes_per_rdma_token + hidden_bytes + - sizeof(SourceMeta)) + - topk_idx); - }; - combine_token( - expected_head >= 0, expected_head, lane_id, hidden_int4, num_topk, combined_x + token_idx * hidden_int4, - combined_topk_weights + token_idx * num_topk, num_max_rdma_chunked_recv_tokens, recv_fn, recv_tw_fn); - } - - // Retired - __syncwarp(); - if (lane_id == 0) rdma_receiver_retired[warp_id] = true; - } else { - // Coordinator - // Sync shared memory status - is_rdma_receiver_sm ? sync_rdma_receiver_smem() : sync_forwarder_smem(); - const auto num_warps_per_rdma_rank = kNumForwarders / kNumRDMARanks; - - int last_rdma_head = 0; - int last_nvl_head[kNumRDMARanks] = {0}; - int dst_rdma_rank = lane_id < kNumRDMARanks ? lane_id : 0; - int dst_nvl_rank = lane_id < NUM_MAX_NVL_PEERS ? lane_id : 0; - EP_STATIC_ASSERT(kNumCombineForwarderWarps <= 32, "Invalid number of forwarder warps"); - while (true) { - // Retired - if (is_rdma_receiver_sm and - __all_sync(0xffffffff, lane_id >= kNumRDMAReceivers or rdma_receiver_retired[lane_id])) - break; - if (not is_rdma_receiver_sm and __all_sync(0xffffffff, lane_id >= kNumForwarders or forwarder_retired[lane_id])) - break; - - // Find minimum head for RDMA ranks - if (is_rdma_receiver_sm) { - int min_head = std::numeric_limits::max(); -#pragma unroll - for (int i = 0; i < kNumRDMAReceivers; ++i) - if (not rdma_receiver_retired[i]) min_head = min(min_head, rdma_receiver_rdma_head[i][dst_rdma_rank]); - if (min_head != std::numeric_limits::max() and - min_head >= last_rdma_head + num_max_rdma_chunked_send_tokens and lane_id < kNumRDMARanks) { - // Phase 4: head feedback path. - // cross-node → fabric-VA st.release on peer's head slot - // self-loop → local atomicAdd - if (peer_rdma_bases != nullptr && dst_rdma_rank != rdma_rank) { - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const uintptr_t my_head_off = reinterpret_cast(rdma_channel_head.buffer(rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - uint64_t* peer_head = reinterpret_cast( - reinterpret_cast(peer_rdma_bases[dst_rank_global]) + my_head_off); - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(peer_head), "l"((uint64_t)min_head) : "memory"); - } else if (dst_rdma_rank == rdma_rank) { - mscclpp::atomicFetchAdd(static_cast(rdma_channel_head.buffer(rdma_rank)), - (uint64_t)(min_head - last_rdma_head), mscclpp::memoryOrderRelease); - } else { - auto dst_offset = rdma_rank * sizeof(uint64_t) + head_send_offset; - auto port_channel_idx = kLowLatencyMode - ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - // Absolute-value RDMA WRITE replaces broken HW atomicAdd. - *rdma_channel_head_send_src.buffer(dst_rdma_rank) = (uint64_t)min_head; - __threadfence_system(); - const auto src_off_head = reinterpret_cast(rdma_channel_head_send_src.buffer(dst_rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - handle.put(dst_offset, src_off_head, sizeof(uint64_t)); - } - last_rdma_head = min_head; - } - } else { -// Find minimum head for NVL ranks -#pragma unroll - for (int i = 0; i < kNumRDMARanks; ++i) { - int min_head = std::numeric_limits::max(); -#pragma unroll - for (int j = 0; j < num_warps_per_rdma_rank; ++j) - if (not forwarder_retired[i * num_warps_per_rdma_rank + j]) - min_head = min(min_head, forwarder_nvl_head[i * num_warps_per_rdma_rank + j][dst_nvl_rank]); - if (min_head != std::numeric_limits::max() and min_head > last_nvl_head[i] and - lane_id < NUM_MAX_NVL_PEERS) - st_relaxed_sys_global(nvl_channel_head.buffer_by(dst_nvl_rank) + i, last_nvl_head[i] = min_head); - } - } - - // Nanosleep and let other warps work - __nanosleep(NUM_WAIT_NANOSECONDS); - } - } - } -} - -void combine(cudaDataType_t type, void* combined_x, float* combined_topk_weights, const bool* is_combined_token_in_rank, - const void* x, const float* topk_weights, const int* combined_rdma_head, const int* combined_nvl_head, - const void* src_meta, const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, - const int* gbl_channel_prefix_matrix, int num_tokens, int num_combined_tokens, int hidden, int num_topk, - void* rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens, - void** buffer_ptrs, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank, - int num_ranks, cudaStream_t stream, int num_channels, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_head_mc, void* nvls_head_dev, - void* nvls_tail_mc, void* nvls_tail_dev, void* const* peer_rdma_bases, void* const* recv_pool_global_ptrs, - const int* ep_combine_recv_idx) { - constexpr int kNumCombineForwarderWarps = 16; - -#define COMBINE_LAUNCH_CASE(num_rdma_ranks) \ - { \ - auto combine_func = low_latency_mode ? combine \ - : combine; \ - LAUNCH_KERNEL(&cfg, combine_func, reinterpret_cast(combined_x), combined_topk_weights, \ - is_combined_token_in_rank, reinterpret_cast(x), topk_weights, combined_rdma_head, \ - combined_nvl_head, reinterpret_cast(src_meta), rdma_channel_prefix_matrix, \ - rdma_rank_prefix_sum, gbl_channel_prefix_matrix, num_tokens, num_combined_tokens, hidden, num_topk, \ - rdma_buffer_ptr, num_max_rdma_chunked_send_tokens, num_max_rdma_chunked_recv_tokens, buffer_ptrs, \ - num_max_nvl_chunked_send_tokens, num_max_nvl_chunked_recv_tokens, rank, num_ranks, \ - port_channel_handles, memory_channel_handles, nvls_head_mc, nvls_head_dev, nvls_tail_mc, \ - nvls_tail_dev, peer_rdma_bases, recv_pool_global_ptrs, ep_combine_recv_idx); \ - } \ - break - - int num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS; - auto num_warps_per_forwarder = std::max(kNumCombineForwarderWarps / num_rdma_ranks, 1); - int num_forwarder_warps = num_rdma_ranks * num_warps_per_forwarder; - EP_HOST_ASSERT(num_forwarder_warps > 0 and num_forwarder_warps % num_rdma_ranks == 0); - EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0); - EP_HOST_ASSERT(num_max_nvl_chunked_recv_tokens / num_rdma_ranks > - std::max(num_max_rdma_chunked_send_tokens, num_max_nvl_chunked_send_tokens)); - EP_HOST_ASSERT(type == CUDA_R_16BF); - -#ifdef EP_DISPATCH_NCCLEP - // FLAT direct-gather combine path. recv_pool_global_ptrs + ep_combine_recv_idx are non-null - // only on the flat direct path (MSCCLPP_EP_DIRECT); every other configuration falls through - // to the unified 2-hop combine below, byte-for-byte unchanged. - if (recv_pool_global_ptrs != nullptr and ep_combine_recv_idx != nullptr) { - // Default: TMA-staged gather (combine_flat_gather_tma). The async copy engine hides the - // remote-NVLink read latency at low register cost, so it wins at every channel count (no - // crossover) and every node scale, and beats NCCL-EP. Escape hatch MSCCLPP_EP_COMBINE_TMA=0 - // falls back to the synchronous lean gather. - static const bool ep_combine_tma_disabled = [] { - const char* e = std::getenv("MSCCLPP_EP_COMBINE_TMA"); - return e != nullptr and std::atoi(e) == 0; - }(); - if (not ep_combine_tma_disabled) { - constexpr int kCmbTmaChunkInt4 = EP_CMB_TMA_CHUNK_INT4; - constexpr int kCmbTmaStages = EP_CMB_TMA_STAGES; - constexpr int kCmbTmaMaxContrib = 8; - // Channel-adaptive warp (token) count: WIDE up to EP_CMB_TMA_WARPS_MAXCH channels (low SM, - // latency-bound), NARROW above it (high SM, where the marginal warp costs more scheduling - // than it buys). Each branch instantiates its own kernel + sets its own SMEM attribute. - const bool cmb_wide = (num_channels <= EP_CMB_TMA_WARPS_MAXCH); -#define COMBINE_FLAT_GATHER_TMA_LAUNCH(num_rdma_ranks, WARPS) \ - { \ - auto tma_func = combine_flat_gather_tma; \ - const size_t cmb_tma_smem = \ - static_cast(WARPS) * kCmbTmaStages * kCmbTmaMaxContrib * kCmbTmaChunkInt4 * sizeof(int4) + \ - static_cast(WARPS) * kCmbTmaStages * sizeof(uint64_t); \ - CUDA_CHECK( \ - cudaFuncSetAttribute(tma_func, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(cmb_tma_smem))); \ - cudaLaunchConfig_t cfg = {static_cast(num_channels * 2), \ - static_cast((WARPS) * 32), \ - cmb_tma_smem, \ - stream, \ - nullptr, \ - 0}; \ - cudaLaunchAttribute a[1]; \ - a[0].id = cudaLaunchAttributeCooperative; \ - a[0].val.cooperative = 1; \ - cfg.attrs = a; \ - cfg.numAttrs = 1; \ - LAUNCH_KERNEL(&cfg, tma_func, reinterpret_cast(combined_x), combined_topk_weights, \ - is_combined_token_in_rank, num_combined_tokens, hidden, num_topk, num_ranks, recv_pool_global_ptrs, \ - ep_combine_recv_idx); \ - } -#define COMBINE_FLAT_GATHER_TMA_CASE(num_rdma_ranks) \ - if (cmb_wide) \ - COMBINE_FLAT_GATHER_TMA_LAUNCH(num_rdma_ranks, EP_CMB_TMA_WARPS_WIDE) \ - else \ - COMBINE_FLAT_GATHER_TMA_LAUNCH(num_rdma_ranks, EP_CMB_TMA_WARPS_NARROW) \ - break - SWITCH_RDMA_RANKS(COMBINE_FLAT_GATHER_TMA_CASE); -#undef COMBINE_FLAT_GATHER_TMA_CASE -#undef COMBINE_FLAT_GATHER_TMA_LAUNCH - return; - } - // Fallback (MSCCLPP_EP_COMBINE_TMA=0): synchronous lean gather. Its full 1024-thread blocks - // hide latency with register MLP + occupancy, which wins when blocks are few but regresses - // at high SM -- so it is gated to a low channel count; above it, the unified kernel's flat - // gather branch runs. - constexpr int kEpCombineLeanMaxChannels = 14; - if (num_channels <= kEpCombineLeanMaxChannels) { -#define COMBINE_FLAT_GATHER_CASE(num_rdma_ranks) \ - { \ - auto gather_func = combine_flat_gather; \ - LAUNCH_KERNEL(&cfg, gather_func, reinterpret_cast(combined_x), combined_topk_weights, \ - is_combined_token_in_rank, num_combined_tokens, hidden, num_topk, num_ranks, recv_pool_global_ptrs, \ - ep_combine_recv_idx); \ - } \ - break - SETUP_LAUNCH_CONFIG(num_channels * 2, 1024, stream); - SWITCH_RDMA_RANKS(COMBINE_FLAT_GATHER_CASE); -#undef COMBINE_FLAT_GATHER_CASE - return; - } - } -#endif - - SETUP_LAUNCH_CONFIG(num_channels * 2, (NUM_MAX_NVL_PEERS + num_forwarder_warps + 1) * 32, stream); - SWITCH_RDMA_RANKS(COMBINE_LAUNCH_CASE); -#undef COMBINE_LAUNCH_CASE -} - -} // namespace internode - -} // namespace ep -} // namespace mscclpp diff --git a/src/ext/ep/ht/kernels/internode_ncclep.cuh b/src/ext/ep/ht/kernels/internode_ncclep.cuh deleted file mode 100644 index 67a27887..00000000 --- a/src/ext/ep/ht/kernels/internode_ncclep.cuh +++ /dev/null @@ -1,1449 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. -// -// NCCL-EP-ported warp-specialized HT dispatch kernel (dispatch_ncclep). -// Guarded by EP_DISPATCH_NCCLEP. Included from internode.cu INSIDE the -// mscclpp::ep::internode namespace, AFTER all shared helpers/buffers are -// defined. When the guard is OFF this file is not compiled and the -// production dispatch<> kernel is used unchanged. -#ifndef MSCCLPP_EP_INTERNODE_NCCLEP_CUH_ -#define MSCCLPP_EP_INTERNODE_NCCLEP_CUH_ -#ifdef EP_DISPATCH_NCCLEP - -// DIAGNOSTIC PROBE (increment-3 de-risk): when 1, the NVL receiver keeps all -// control flow (index/head/tail) but SKIPS the actual data copies. This makes -// recv_x contents wrong for cross-GPU tokens (combine FAILs) but measures the -// dispatch-time UPPER BOUND of eliminating the cross-GPU receiver drain (the -// payoff ceiling of the full cross-GPU peer-map direct-write rework). Set to 0 -// for the real kernel. -#ifndef EP_NCCLEP_DRAIN_NOOP -#define EP_NCCLEP_DRAIN_NOOP 0 -#endif - -// Increment 4 (VMM pool + TMA): when 1, the cross-GPU forwarder writes hidden to -// the destination peer's VMM recv-output pool via cp.async.bulk (TMA) instead of -// UNROLLED_WARP_COPY. Requires recv_pool_ptrs to be cuMem (FABRIC/POSIX-FD) -// peer VAs (gpuCallocPhysical), NOT cudaIpc-mapped memory (which crashes TMA). -#ifndef EP_NCCLEP_TMA -#define EP_NCCLEP_TMA 0 -#endif - -// B-depth-3 (full-token variant): per-sender-warp TMA S2G ring geometry (shared with the -// host launcher so it can size the DYNAMIC shared-memory allocation + cudaFuncSetAttribute -// opt-in). WHOLE-token (14336B) tiles -> just 1 cp.async.bulk S2G per destination per token -// (matches NCCL-EP's single whole-token op), vs 2 for the half-token / 14 for the old 1024B -// sub-chunk. The ring is kNumDispatchRDMASenderWarps(6) x NSTAGE(2) x 14336 = 172032 B (168 KB) -// -- full-token forces NSTAGE=2 to stay under the GB200 ~227KB/block opt-in cap (NSTAGE>=3 would -// be >=252KB). chunk MUST be a multiple of 128 (TMA align); 14336 = the full bf16 hidden so each -// token is a single chunk (nchunks=1), and the pipeline overlap is purely cross-token. -#ifndef EP_TMA_SND_CHUNK_BYTES -#define EP_TMA_SND_CHUNK_BYTES 14336 -#endif -#ifndef EP_TMA_SND_NSTAGE -#define EP_TMA_SND_NSTAGE 2 -#endif - -// Increment 5 (inc5): runtime gate for the SENDER direct-write path. When set -// (env MSCCLPP_EP_DIRECT=1), kRDMASender writes each token's hidden straight to -// the destination GPU's domain-wide recv pool (recv_pool_global_ptrs[dst_global]), -// skipping the rdma_channel hidden bounce + forwarder hidden transpose; metadata -// (scales/topk/meta) still flows the legacy ring->forwarder->receiver path. -__constant__ int kEpDirect; - -// Increment 6 (inc6): runtime gate for the FLAT all-sender dispatch path. When -// set (env MSCCLPP_EP_FLAT=1, requires kEpDirect), the sender ALSO writes each -// token's metadata (SourceMeta + scales + topk) straight into the destination -// pool's meta region at the token's final recv slot, a per-(dst,channel) fabric -// counter signals completion, and a lightweight LOCAL consumer copies the pool -// meta into the recv_* output tensors -- eliminating the forwarder + both -// coordinator roles so dispatch can launch an all-sender grid (num_channels = -// num_sms). Unset => byte-identical inc5 2-hop metadata path. -__constant__ int kEpFlat; - -template -__global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS) * 32), 1) - dispatch_ncclep(int4* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv_topk_weights, - SourceMeta* recv_src_meta, const int4* x, const float* x_scales, const int64_t* topk_idx, - const float* topk_weights, int* send_rdma_head, int* send_nvl_head, - int* recv_rdma_channel_prefix_matrix, int* recv_gbl_channel_prefix_matrix, - const int* rdma_channel_prefix_matrix, const int* recv_rdma_rank_prefix_sum, - const int* gbl_channel_prefix_matrix, const int* recv_gbl_rank_prefix_sum, int num_tokens, - int hidden_int4, int num_scales, int num_topk, int num_experts, const bool* is_token_in_rank, - void* rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens, - void** buffer_ptrs, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, - int rank, int num_ranks, mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, - // Phase 3 NVLS counter pointers (nullptr → fall back to PortChannel/atomicAdd path). - void* nvls_head_mc, void* nvls_head_dev, void* nvls_tail_mc, void* nvls_tail_dev, - // Phase 4: per-peer fabric-IPC base pointers; when non-null, cross-node data - // PUTs go directly through NVL72 fabric VA instead of `handle.put` over IB. - void* const* peer_rdma_bases, - // Increment 4: per-peer base pointers of the VMM-allocated (cuMem FABRIC/POSIX-FD) - // recv-output pool. recv_pool_ptrs[peer] points at peer's pool header; non-null - // enables the cross-GPU forwarder direct-write of hidden into the destination's - // final recv_x (TMA-eligible peer VA); nullptr = legacy receiver-drain path. - void* const* recv_pool_ptrs, - // Increment 5 (inc5): domain-wide recv-pool bases indexed by GLOBAL rank - // (all num_ranks). Non-null + kEpDirect => sender writes hidden direct here. - void* const* recv_pool_global_ptrs, - // Increment 5 combine-direct (Stage 1): [num_tokens * num_ranks] gather map; - // sender writes recv_idx per (token, dst global rank) for combine to gather. - int* ep_combine_recv_idx) { - enum class WarpRole { - kRDMASender, - kRDMASenderCoordinator, - kRDMAAndNVLForwarder, - kForwarderCoordinator, - kNVLReceivers - }; - - const auto sm_id = static_cast(blockIdx.x); - const auto thread_id = static_cast(threadIdx.x), warp_id = thread_id / 32, lane_id = get_lane_id(); - // inc6 (kEpFlat): the FLAT all-sender path launches 1 block per channel - // (gridDim.x == num_channels) with no forwarder blocks, so every block is a - // sender. The inc5 2-hop path launches 2 blocks per channel (sender + forwarder), - // hence the /2; channel_id and is_forwarder follow the same split. - const auto num_channels = kEpFlat ? static_cast(gridDim.x) : static_cast(gridDim.x) / 2; - const auto channel_id = kEpFlat ? sm_id : sm_id / 2; - const bool is_forwarder = kEpFlat ? false : (sm_id % 2 == 0); - // inc7 (atomicthr): per-channel routing-lock removal is a CROSSOVER vs channel - // count. num_channels == config.num_sms/2 (buffer.cc), and tokens/channel = - // num_tokens/num_channels. At LOW channel count (many tokens/channel) the - // per-token sequential lock dominates -> atomicAdd slot assignment (no spin) - // wins big: 8 ch (NSM16) -21.6%, 12 ch (NSM24) -21.1%, 14 ch (NSM28) -4.2%. - // At HIGH channel count (few tokens/channel) the lock was already cheap AND the - // 6 sender warps contend on the shared-mem atomics -> net regression: 16 ch - // (NSM32) +8.6%. Measured crossover is between 14 and 16 channels (~14.7), so - // use atomics at <=14 channels (every validated win) and keep the proven lock - // above it (the validated regression). num_channels is grid-uniform so this is - // a block-uniform branch (no warp divergence). - constexpr int kEpAtomicRouteMaxChannels = 14; - const bool ep_use_atomic_route = kEpFlat and (num_channels <= kEpAtomicRouteMaxChannels); - const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS; - - const auto role_meta = [=]() -> std::pair { - if (is_forwarder) { - if (warp_id < NUM_MAX_NVL_PEERS) { - return {WarpRole::kRDMAAndNVLForwarder, (warp_id + channel_id) % NUM_MAX_NVL_PEERS}; - } else { - return {WarpRole::kForwarderCoordinator, warp_id - NUM_MAX_NVL_PEERS}; - } - } else if (warp_id < kNumDispatchRDMASenderWarps) { - return {WarpRole::kRDMASender, -1}; - } else if (warp_id == kNumDispatchRDMASenderWarps) { - return {WarpRole::kRDMASenderCoordinator, -1}; - } else { - return {WarpRole::kNVLReceivers, (warp_id + channel_id - kNumDispatchRDMASenderWarps) % NUM_MAX_NVL_PEERS}; - } - }(); - auto warp_role = role_meta.first; - auto target_rank = role_meta.second; // Not applicable for RDMA senders - EP_DEVICE_ASSERT(num_warps == kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NVL_PEERS); - if (thread_id == 0 && sm_id == 0) { - (void)0; - } - - // Data checks - EP_DEVICE_ASSERT(num_topk <= 32); - - // RDMA symmetric layout (packed-bool size guard is at namespace scope via NvlPackT). - // Snapshot the original base before SymBuffer constructors advance it; used - // below to compute MR-relative offsets for handle.put(). - void* const rdma_buffer_ptr_base = rdma_buffer_ptr; - auto hidden_bytes = hidden_int4 * sizeof(int4); - // inc5: under kEpDirect the rdma ring carries NO hidden (it goes direct to the - // dest pool), so the per-token slot excludes hidden -> the coordinator RDMAs - // only metadata. hidden_bytes (above) stays full for the pool/source copies. - const int ring_hidden_int4 = kEpDirect ? 0 : hidden_int4; - const auto ring_hidden_bytes = static_cast(ring_hidden_int4) * sizeof(int4); - auto num_bytes_per_rdma_token = get_num_bytes_per_rdma_token(ring_hidden_int4, num_scales, num_topk, num_topk); - auto rdma_channel_data = - SymBuffer(rdma_buffer_ptr, num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token, kNumRDMARanks, - channel_id, num_channels); - auto rdma_channel_meta = - SymBuffer(rdma_buffer_ptr, NUM_MAX_NVL_PEERS * 2 + 2, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_head = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_tail = SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - // Scratch slots (one uint64_t per (channel, peer) per kind) holding the - // absolute counter values used as RDMA WRITE source. Replaces the broken - // HW atomicAdd path on Azure CX-7 RoCE (IBV_ATOMIC_NONE) — each tail/head - // slot has a single writer per peer, so atomicity is unnecessary; an - // absolute-value RDMA WRITE through the same QP as the data PUT preserves - // ordering by IB semantics. - auto rdma_channel_tail_send_src = - SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - auto rdma_channel_head_send_src = - SymBuffer(rdma_buffer_ptr, 1, kNumRDMARanks, channel_id, num_channels); - - auto data_send_offset = - sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * kNumRDMARanks * channel_id; - auto data_recv_offset = sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * - kNumRDMARanks * (channel_id + num_channels); - auto meta_offset = - sizeof(int8_t) * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) * kNumRDMARanks * num_channels * 2; - auto meta_send_offset = meta_offset + sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2) * kNumRDMARanks * channel_id; - auto meta_recv_offset = - meta_offset + sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2) * kNumRDMARanks * (channel_id + num_channels); - auto head_offset = meta_offset + sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2) * kNumRDMARanks * num_channels * 2; - auto head_send_offset = head_offset + sizeof(uint64_t) * kNumRDMARanks * channel_id; - auto tail_offset = head_offset + sizeof(uint64_t) * kNumRDMARanks * num_channels; - auto tail_send_offset = tail_offset + sizeof(uint64_t) * kNumRDMARanks * channel_id; - - // NVL buffer layouts - // NOTES: `rs_wr_buffer_ptr` means "Read for Senders, Write for Receivers", `ws_rr_buffer_ptr` means "Write for - // Senders, Read for Receivers" - void *rs_wr_buffer_ptr = nullptr, *ws_rr_buffer_ptr = nullptr; - int rs_wr_rank = 0, ws_rr_rank = 0; - if (warp_role == WarpRole::kRDMAAndNVLForwarder) - rs_wr_buffer_ptr = buffer_ptrs[nvl_rank], ws_rr_buffer_ptr = buffer_ptrs[target_rank], rs_wr_rank = nvl_rank, - ws_rr_rank = target_rank; - if (warp_role == WarpRole::kNVLReceivers) - rs_wr_buffer_ptr = buffer_ptrs[target_rank], ws_rr_buffer_ptr = buffer_ptrs[nvl_rank], rs_wr_rank = target_rank, - ws_rr_rank = nvl_rank; - - // Allocate buffers - auto nvl_channel_x = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * hidden_int4, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_src_meta = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_x_scales = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_scales, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_topk_idx = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_topk_weights = AsymBuffer(ws_rr_buffer_ptr, num_max_nvl_chunked_recv_tokens * num_topk, - NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_prefix_start = - AsymBuffer(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_prefix_end = - AsymBuffer(ws_rr_buffer_ptr, kNumRDMARanks, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - auto nvl_channel_head = AsymBuffer(rs_wr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id, num_channels, ws_rr_rank) - .advance_also(ws_rr_buffer_ptr); - auto nvl_channel_tail = AsymBuffer(ws_rr_buffer_ptr, 1, NUM_MAX_NVL_PEERS, channel_id, num_channels, rs_wr_rank) - .advance_also(rs_wr_buffer_ptr); - - // RDMA sender warp synchronization - __shared__ volatile int rdma_send_next_token_idx; - __shared__ volatile int rdma_send_channel_tail[kNumRDMARanks]; - __shared__ volatile int rdma_send_channel_next_tail[kNumRDMARanks]; - // inc5: per-dst-GPU (global rank) final recv_x base index + running count, used - // by kRDMASender under kEpDirect to write hidden straight to the dest pool. - __shared__ int ep_base[kNumRDMARanks * NUM_MAX_NVL_PEERS]; - __shared__ int ep_count[kNumRDMARanks * NUM_MAX_NVL_PEERS]; - auto sync_rdma_sender_smem = []() { asm volatile("bar.sync 0, %0;" ::"r"((kNumDispatchRDMASenderWarps + 1) * 32)); }; - - // Forward warp synchronization - __shared__ volatile int forward_channel_head[NUM_MAX_NVL_PEERS][kNumRDMARanks]; - __shared__ volatile bool forward_channel_retired[NUM_MAX_NVL_PEERS]; - auto sync_forwarder_smem = []() { asm volatile("bar.sync 1, %0;" ::"r"((NUM_MAX_NVL_PEERS + 1) * 32)); }; - -#if EP_NCCLEP_TMA - // Increment 4 (TMA): per-forwarder-warp SMEM staging tile + mbarrier for the - // cross-GPU direct-write of hidden into the destination peer's VMM recv pool. - // One tile/mbar per forwarder warp (warp_id 0..NUM_MAX_NVL_PEERS-1); the - // per-token hidden is streamed HBM->SMEM->peer VA in kEpTmaTileBytes chunks. - static constexpr int kEpTmaTileBytes = 4096; - __shared__ alignas(128) uint8_t ep_tma_tile[NUM_MAX_NVL_PEERS][kEpTmaTileBytes]; - __shared__ alignas(8) uint64_t ep_tma_mbar[NUM_MAX_NVL_PEERS]; - // inc5 sender-direct TMA ring (NCCL-EP g2s_s2g_pipe pattern): per-sender-warp ring - // of kEpTmaSndNStage tiles+mbars. Stage A (inc6, pipelined): the ring is now drained - // LAZILY with a register-tracked in-flight bound (NCCL-EP hybrid_ep.cuh model) so the - // S2G NVLink writes of consecutive chunks overlap with each other AND with the next - // chunk's G2S HBM load, instead of the old per-chunk full `wait_group 0` drain (which - // serialized G2S vs S2G and exposed the full NVLink round-trip at every chunk -- the - // 8-node low-SM dispatch cliff). Pipeline: produce 1 chunk ahead, keep kEpTmaSndInFlight - // S2G groups in flight, reuse a tile only after its S2G has drained (proven safe for - // produce-ahead=1 + InFlight=NStage-2; the drained-stage index trails the refill stage - // by exactly one iteration). B-depth-3: chunk is now HALF-token (7168B) and the ring - // lives in DYNAMIC shared memory (6*4*7168 = 168KB > 48KB static cap), so each token is - // 2 cp.async.bulk S2G per destination instead of 14 -- fewer, larger TMA descriptors - // (NCCL-EP issues 1 whole-token op; this halves the gap). Host opts in to the >48KB ring - // via cudaFuncAttributeMaxDynamicSharedMemorySize + cfg.dynamicSmemBytes (EP_TMA_SND_*). - static constexpr int kEpTmaSndChunkBytes = EP_TMA_SND_CHUNK_BYTES; - static constexpr int kEpTmaSndNStage = EP_TMA_SND_NSTAGE; // ring depth (lazy drain; >=3 for within-token overlap) - // S2G groups kept in flight: NStage-2 when there is within-token slack (half-token NSTAGE>=4), - // clamped to >=1 so the full-token NSTAGE=2 case keeps 1 cross-token S2G in flight. - static constexpr int kEpTmaSndInFlight = (kEpTmaSndNStage - 2) >= 1 ? (kEpTmaSndNStage - 2) : 1; - EP_STATIC_ASSERT(kEpTmaSndInFlight >= 1, "InFlight must be >=1 for pipelining"); - EP_STATIC_ASSERT(kEpTmaSndInFlight < kEpTmaSndNStage, "InFlight must be < NStage (drain-before-reuse)"); - EP_STATIC_ASSERT(kEpTmaSndChunkBytes % 128 == 0, "TMA tile must be 128B aligned"); - // Dynamic per-warp ring: ep_tma_tile_snd_dyn[(warp_id*NStage + stage) * chunk]. 128B-aligned - // base for TMA. The mbarrier array stays static (tiny: 6*4*8 = 192B). - extern __shared__ __align__(128) uint8_t ep_tma_tile_snd_dyn[]; - __shared__ alignas(8) uint64_t ep_tma_mbar_snd[kNumDispatchRDMASenderWarps][kEpTmaSndNStage]; -#endif - - if (warp_role == WarpRole::kRDMASender) { - // Get tasks - int token_start_idx, token_end_idx; - get_channel_task_range(num_tokens, num_channels, channel_id, token_start_idx, token_end_idx); - - // Clean shared memory - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA ranks"); - (warp_id == 0 and lane_id == 0) ? (rdma_send_next_token_idx = token_start_idx) : 0; - (warp_id == 0 and lane_id < kNumRDMARanks) ? (rdma_send_channel_tail[lane_id] = 0) : 0; - (warp_id == 0 and lane_id < kNumRDMARanks) ? (rdma_send_channel_next_tail[lane_id] = 0) : 0; - - // Send number of tokens in this channel by `-value - 1` - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS * 2 + 2 <= 32, "Invalid number of NVL peers"); - for (int dst_rdma_rank = warp_id; dst_rdma_rank < kNumRDMARanks; dst_rdma_rank += kNumDispatchRDMASenderWarps) { - auto dst_ptr = dst_rdma_rank == rdma_rank ? rdma_channel_meta.recv_buffer(dst_rdma_rank) - : rdma_channel_meta.send_buffer(dst_rdma_rank); - if (lane_id < NUM_MAX_NVL_PEERS) { - dst_ptr[lane_id] = - -(channel_id == 0 ? 0 - : gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id) * num_channels + - channel_id - 1]) - - 1; - } else if (lane_id < NUM_MAX_NVL_PEERS * 2) { - dst_ptr[lane_id] = - -gbl_channel_prefix_matrix[(dst_rdma_rank * NUM_MAX_NVL_PEERS + lane_id - NUM_MAX_NVL_PEERS) * - num_channels + - channel_id] - - 1; - } else if (lane_id == NUM_MAX_NVL_PEERS * 2) { - dst_ptr[lane_id] = - -(channel_id == 0 ? 0 : rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id - 1]) - 1; - } else if (lane_id == NUM_MAX_NVL_PEERS * 2 + 1) { - dst_ptr[lane_id] = -rdma_channel_prefix_matrix[dst_rdma_rank * num_channels + channel_id] - 1; - } - __syncwarp(); - - if (dst_rdma_rank == rdma_rank) continue; - - // Issue RDMA for non-local ranks - if (peer_rdma_bases != nullptr) { - // Phase 4: deliver meta directly via NVL72 fabric VA (cooperative - // int copy across the warp). Replaces port_channel.put which is - // unreliable cross-node on Azure CX-7 RoCE without flush. - const auto num_bytes = sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2); - const auto dst_offset = rdma_rank * num_bytes + meta_recv_offset; - const auto src_offset = dst_rdma_rank * num_bytes + meta_send_offset; - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const int* src_p = reinterpret_cast(reinterpret_cast(rdma_buffer_ptr_base) + src_offset); - int* dst_p = reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + dst_offset); - const int n_int = (int)(num_bytes / sizeof(int)); - for (int k = lane_id; k < n_int; k += 32) { - dst_p[k] = src_p[k]; - } - __threadfence_system(); - } else if (lane_id == 0) { - auto num_bytes = sizeof(int) * (NUM_MAX_NVL_PEERS * 2 + 2); - auto dst_offset = rdma_rank * num_bytes + meta_recv_offset; - auto src_offset = dst_rdma_rank * num_bytes + meta_send_offset; - auto port_channel_idx = kLowLatencyMode - ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - port_channel_handles[port_channel_idx].put(dst_offset, src_offset, num_bytes); - // port_channel_handles[port_channel_idx].flush(); - } - __syncwarp(); - } - - // inc5: precompute per-dst-GPU base index into the destination's recv_x pool. - // base[dg] = dst's recv_gbl_rank_prefix_sum[rank-1] (peer-read from dg's pool - // header) + this channel's start offset for (rank -> dg). Running count is added - // per token in the seq-locked section below. No extra barrier: covered by the - // sync_rdma_sender_smem() right after. - const int64_t ep_pool_header_bytes = ((static_cast(num_ranks) * 4 + 127) / 128) * 128; - // inc6 (kEpFlat): byte offset of the per-token metadata region in the recv pool, - // and the per-token meta slot size. MUST match Config::get_recv_pool_meta_base / - // recv_pool_bytes_static (config.hpp): header + worst-case hidden region, padded to - // 128; kEpFlatMetaBytes per token. The sender writes SourceMeta + scales + topk here - // so a local consumer can produce recv_* without the forwarder/ring 2-hop. - constexpr int64_t kEpFlatPoolMaxTokens = 65536; // == Config::kEpRecvPoolMaxTokens - constexpr int64_t kEpFlatPoolMaxHiddenBytes = 16384; // == Config::kEpRecvPoolMaxHiddenBytes - constexpr int64_t kEpFlatMetaBytes = 128; // == Config::kEpRecvPoolMetaBytes - const int64_t ep_meta_base = - ((ep_pool_header_bytes + kEpFlatPoolMaxTokens * kEpFlatPoolMaxHiddenBytes + 127) / 128) * 128; - if (kEpDirect and recv_pool_global_ptrs != nullptr) { - for (int dg = warp_id * 32 + lane_id; dg < num_ranks; dg += kNumDispatchRDMASenderWarps * 32) { - const int* dst_prefix = reinterpret_cast(recv_pool_global_ptrs[dg]); - int gbl_base = (rank > 0) ? static_cast(ld_volatile_global(dst_prefix + (rank - 1))) : 0; - int ch_start = (channel_id == 0) ? 0 : gbl_channel_prefix_matrix[dg * num_channels + channel_id - 1]; - ep_base[dg] = gbl_base + ch_start; - ep_count[dg] = 0; - } - } - sync_rdma_sender_smem(); - - // Iterate over tokens and copy into buffer - int64_t token_idx; - int cached_rdma_channel_head = 0, last_rdma_tail_idx = -1; - auto send_buffer = - lane_id == rdma_rank ? rdma_channel_data.recv_buffer(lane_id) : rdma_channel_data.send_buffer(lane_id); -#if EP_NCCLEP_TMA - // B-depth-1: TMA S2G pipeline state carried ACROSS tokens so the per-token full drain - // (wait_group 0) is eliminated -- up to kEpTmaSndInFlight S2G groups stay in flight across - // token boundaries. ep_tma_gbase = global ring-stage base (continuous chunk index, advances - // by nchunks each token); ep_tma_in_flight = S2G groups committed but not yet drained. - // Safe across tokens for the same reason as within a token: NStage(4) > InFlight(2), so the - // stage a new G2S reuses (g) was last used by chunk g-NStage <= g-4, which is already drained - // by the lazy wait_group.read(InFlight) before it is overwritten. - int ep_tma_gbase = 0; - int ep_tma_in_flight = 0; -#endif - for (token_idx = token_start_idx + warp_id; token_idx < token_end_idx; token_idx += kNumDispatchRDMASenderWarps) { - // Read RDMA rank existence - NvlPackT is_token_in_rank_uint64 = 0; - if (lane_id < kNumRDMARanks) - is_token_in_rank_uint64 = - *reinterpret_cast(is_token_in_rank + token_idx * num_ranks + lane_id * NUM_MAX_NVL_PEERS); - - // inc7 (atomicthr): under kEpFlat AND low channel count the per-channel - // SEQUENTIAL LOCK is REMOVED (ep_use_atomic_route). Slot assignment uses - // atomicAdd on the shared counters (below) -> dense but order-independent - // packing. Correctness holds because (a) the receiver's recv_x check is - // amin/amax over each source's cross-rank block (order-agnostic; every token - // from a given source lands in that source's reserved [gbl_base, ..) range - // regardless of within-channel order), and (b) combine gathers via the self- - // consistent ep_combine_recv_idx map. The lock path (high channel count or - // non-flat 2-hop, whose coordinator forwards a MONOTONIC rdma_send_channel_tail - // to peers) is UNCHANGED -- keeps the spin-wait + plain ++. - while (!ep_use_atomic_route and lane_id == 0 and rdma_send_next_token_idx != token_idx); - __syncwarp(); - - // Acquire next tail (atomic when lock removed: 6 sender warps race; plain ++ - // under the lock where it is single-threaded == the committed behavior). - int rdma_tail_idx = -1; - if (is_token_in_rank_uint64 != 0) { - rdma_tail_idx = ep_use_atomic_route ? atomicAdd(const_cast(rdma_send_channel_next_tail + lane_id), 1) - : rdma_send_channel_next_tail[lane_id]++; - // inc6 (kEpFlat): no forwarder drains the ring, so the head never advances - // and this flow-control wait would deadlock. The all-sender path writes - // hidden + metadata straight to the dest pool; rdma_tail_idx is still used - // below only as the per-node "token in rank" flag (>= 0) and combine head. - if (!kEpFlat) { - while (rdma_tail_idx - cached_rdma_channel_head >= num_max_rdma_chunked_recv_tokens) { - // Phase 4: head feedback path \u2014 cross-node uses fabric-VA store, - // self-loop uses local atomic. Both end up in rdma_channel_head; - // single read via ld_volatile_global covers them. NVLS removed. - cached_rdma_channel_head = static_cast(ld_volatile_global(rdma_channel_head.buffer(lane_id))); - } - } - } - __syncwarp(); - - // inc5: compute this token's final recv_x index for each destination GPU, - // inside the seq-locked region so the running count matches the receiver's - // drain order. lane = source NODE; one writer per ep_count[dg]. - // inc7 (scanexp1): the running-count increments MUST stay ordered (inside the - // lock), but the *global* breadcrumb stores (ep_combine_recv_idx, send_rdma_head) - // target unique per-(token,dst) addresses and are order-independent, so they are - // DEFERRED to after the lock release. This shrinks the serialized critical section - // to just the shared-memory counter math, taking the per-token serialized - // global-store latency out of the lock chain (NSM16 diagnostic). - int ep_my_idx[NUM_MAX_NVL_PEERS]; - if (kEpDirect and lane_id < kNumRDMARanks and rdma_tail_idx >= 0) { - const bool* bvals = reinterpret_cast(&is_token_in_rank_uint64); -#pragma unroll - for (int g = 0; g < NUM_MAX_NVL_PEERS; ++g) - if (bvals[g]) { - const int dg = lane_id * NUM_MAX_NVL_PEERS + g; - // inc7 (atomicthr): atomic running-count when the lock is removed - // (order-independent); plain ++ under the lock (== committed, serialized). - if (ep_use_atomic_route) { - ep_my_idx[g] = ep_base[dg] + atomicAdd(&ep_count[dg], 1); - } else { - ep_my_idx[g] = ep_base[dg] + ep_count[dg]; - ep_count[dg] += 1; - } - } - } - - // Update last token tail. In-loop writes are sequenced by the - // per-channel sequential lock and the warp-stride property of the - // token loop, so monotonicity is guaranteed and a plain - // st_release_cta is correct AND faster than atomicMax (which - // would serialize through L2 if the compiler can't infer shared - // address space). The epilogue (out of the seq-lock contract for - // the highest in-rank slot) needs atomicMax separately. - if (last_rdma_tail_idx >= 0) - st_release_cta(const_cast(rdma_send_channel_tail + lane_id), last_rdma_tail_idx + 1); - last_rdma_tail_idx = rdma_tail_idx; - - // Release sequential lock (lock path only; atomic-route removed the lock above) - if (!ep_use_atomic_route and lane_id == 0) rdma_send_next_token_idx += 1; - - // inc7 (scanexp1): DEFERRED global breadcrumb stores, now OUTSIDE the lock so - // they run in parallel across the sender warps. Addresses are unique per - // (token_idx, dst), so out-of-lock ordering is correctness-preserving. ep_my_idx[g] - // and rdma_tail_idx persist in registers from the in-lock computation above. - if (kEpDirect and ep_combine_recv_idx != nullptr and not kCachedMode and lane_id < kNumRDMARanks and - rdma_tail_idx >= 0) { - const bool* bvals = reinterpret_cast(&is_token_in_rank_uint64); -#pragma unroll - for (int g = 0; g < NUM_MAX_NVL_PEERS; ++g) - if (bvals[g]) - ep_combine_recv_idx[static_cast(token_idx) * num_ranks + lane_id * NUM_MAX_NVL_PEERS + g] = - ep_my_idx[g]; - } - - // Store RDMA head for combine (deferred, post-lock; unique per (token, node)). - if (lane_id < kNumRDMARanks and not kCachedMode) - send_rdma_head[token_idx * kNumRDMARanks + lane_id] = rdma_tail_idx; - - // Broadcast tails - SourceMeta src_meta; - int num_topk_ranks = 0, topk_ranks[kNumTopkRDMARanks]; - void* dst_send_buffers[kNumTopkRDMARanks]; -#pragma unroll - for (int i = 0, slot_idx; i < kNumRDMARanks; ++i) - if ((slot_idx = __shfl_sync(0xffffffff, rdma_tail_idx, i)) >= 0) { - slot_idx = slot_idx % num_max_rdma_chunked_recv_tokens; - topk_ranks[num_topk_ranks] = i; - auto recv_is_token_in_rank_uint64 = warpBroadcast(is_token_in_rank_uint64, i); - auto recv_is_token_in_rank_values = reinterpret_cast(&recv_is_token_in_rank_uint64); - if (lane_id == num_topk_ranks) src_meta = SourceMeta(rdma_rank, recv_is_token_in_rank_values); - dst_send_buffers[num_topk_ranks++] = - reinterpret_cast(warpBroadcast(send_buffer, i)) + slot_idx * num_bytes_per_rdma_token; - } - EP_DEVICE_ASSERT(num_topk_ranks <= kNumTopkRDMARanks); - - // Copy `x` into the send path. inc5 (kEpDirect): write hidden DIRECTLY to - // each destination GPU's recv_x pool at its final index; the per-NODE ring - // hidden region is left unwritten (the forwarder skips it under kEpDirect). - // Else (inc4a): broadcast hidden into the ring slots for forwarder transpose. - if (kEpDirect) { - // inc5: gather every destination GPU's recv_x slot for this token, then - // read `x` from source ONCE and broadcast it to all of them (mirrors the - // inc4a st_broadcast pattern). A naive per-destination UNROLLED_WARP_COPY - // re-reads `x` from HBM for every target GPU, wasting source bandwidth. - int4* ep_dst_pools[kNumTopkRDMARanks * NUM_MAX_NVL_PEERS]; - int ep_num_pools = 0; -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) { - const int node = topk_ranks[j]; - NvlPackT bools_j = warpBroadcast(is_token_in_rank_uint64, node); - const bool* bvals_j = reinterpret_cast(&bools_j); -#pragma unroll - for (int g = 0; g < NUM_MAX_NVL_PEERS; ++g) { - if (not bvals_j[g]) continue; - const int dg = node * NUM_MAX_NVL_PEERS + g; - const int recv_idx = __shfl_sync(0xffffffff, ep_my_idx[g], node); - ep_dst_pools[ep_num_pools++] = - reinterpret_cast(reinterpret_cast(recv_pool_global_ptrs[dg]) + ep_pool_header_bytes + - static_cast(recv_idx) * hidden_bytes); - } - } -#if EP_NCCLEP_TMA - // inc5 + TMA ring, LANE-STRIPED broadcast (NCCL-EP hybrid_ep.cuh S2G shape), - // Stage A PIPELINED: lane 0 streams each chunk HBM->SMEM (G2S, mbarrier - // complete_tx); the whole warp waits that mbarrier then fans the staged tile out - // S2G with ONE destination pool per lane. Unlike the old per-chunk `wait_group 0` - // (which fully drained S2G before the next G2S, serializing the two and exposing - // the full NVLink round-trip every chunk), we now PRODUCE ONE CHUNK AHEAD and drain - // LAZILY: up to kEpTmaSndInFlight S2G groups stay in flight, so consecutive chunks' - // NVLink writes overlap each other and the next chunk's HBM load. A ring tile is - // refilled only one iteration after its S2G drained (proven safe for produce-ahead=1 - // + InFlight=NStage-2). fence.proxy publishes the writes before the dispatch barrier. - if (ep_num_pools > 0) { - const uint8_t* src_b = reinterpret_cast(x + token_idx * hidden_int4); - const size_t total = static_cast(hidden_bytes); - const int nchunks = static_cast((total + kEpTmaSndChunkBytes - 1) / kEpTmaSndChunkBytes); - - // G2S issue for chunk `ci` into its ring stage (caller must guard lane 0 only). - auto issue_g2s = [&](int ci) { - const size_t goff = static_cast(ci) * kEpTmaSndChunkBytes; - const size_t remain = total - goff; - uint32_t csz = - static_cast(remain < (size_t)kEpTmaSndChunkBytes ? remain : (size_t)kEpTmaSndChunkBytes); - csz = (csz + 15u) & ~15u; - const int stage = (ep_tma_gbase + ci) % kEpTmaSndNStage; - const uint32_t smem_tile = static_cast(__cvta_generic_to_shared( - ep_tma_tile_snd_dyn + (warp_id * kEpTmaSndNStage + stage) * kEpTmaSndChunkBytes)); - const uint32_t smem_mbar = - static_cast(__cvta_generic_to_shared(&ep_tma_mbar_snd[warp_id][stage])); - asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_mbar)); - fenceProxyAsyncSharedCta(); - asm volatile( - "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " - "[%0], [%1], %2, [%3];" ::"r"(smem_tile), - "l"(src_b + goff), "r"(csz), "r"(smem_mbar) - : "memory"); - uint64_t st; - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 %0, [%1], %2;" - : "=l"(st) - : "r"(smem_mbar), "r"(csz)); - }; - - // Prologue: prime the first stage (the loop produces one chunk ahead). - if (lane_id == 0) issue_g2s(0); - __syncwarp(); - - for (int c = 0; c < nchunks; ++c) { - // Produce the next chunk's G2S so it overlaps this chunk's S2G drain. - if (lane_id == 0 and c + 1 < nchunks) issue_g2s(c + 1); - __syncwarp(); - - const size_t off = static_cast(c) * kEpTmaSndChunkBytes; - const size_t remain = total - off; - uint32_t csz = - static_cast(remain < (size_t)kEpTmaSndChunkBytes ? remain : (size_t)kEpTmaSndChunkBytes); - csz = (csz + 15u) & ~15u; - const int stage = (ep_tma_gbase + c) % kEpTmaSndNStage; - const uint32_t smem_tile = static_cast(__cvta_generic_to_shared( - ep_tma_tile_snd_dyn + (warp_id * kEpTmaSndNStage + stage) * kEpTmaSndChunkBytes)); - const uint32_t smem_mbar = - static_cast(__cvta_generic_to_shared(&ep_tma_mbar_snd[warp_id][stage])); - - // Whole warp waits on this chunk's G2S mbarrier so every S2G lane sees the tile. - uint32_t done = 0; - while (!done) { - asm volatile( - "{ .reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], 0;" - " selp.u32 %0, 1, 0, p; }" - : "=r"(done) - : "r"(smem_mbar)); - } - fenceProxyAsyncSharedCta(); - // S2G: lane-striped fan-out -- one destination pool per lane. - bool issued = false; - for (int j = lane_id; j < ep_num_pools; j += 32) { - uint8_t* dst_b = reinterpret_cast(ep_dst_pools[j]) + off; - asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst_b), "r"(smem_tile), - "r"(csz) - : "memory"); - issued = true; - } - if (issued) asm volatile("cp.async.bulk.commit_group;"); - ++ep_tma_in_flight; - // Lazy drain: bound in-flight S2G to kEpTmaSndInFlight (compile-time immediate). - // wait_group.read also guarantees the drained group's SMEM reads are done, so - // its ring tile is safe to refill by a future G2S one iteration later. - if (ep_tma_in_flight > kEpTmaSndInFlight) { - if (issued) asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(kEpTmaSndInFlight) : "memory"); - --ep_tma_in_flight; - } - __syncwarp(); - } - // B-depth-1: NO per-token drain -- carry in-flight S2G groups + the ring stage - // (ep_tma_gbase) to the next token so consecutive tokens' S2G NVLink writes overlap - // each other (removes the per-token full-drain bubble = ~1 NVLink RT/token at low SM). - ep_tma_gbase += nchunks; - } - __syncwarp(); -#else - auto st_pool_broadcast = [=](const int key, const int4& value) { -#pragma unroll - for (int j = 0; j < ep_num_pools; ++j) st_na_global(ep_dst_pools[j] + key, value); - }; - UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_pool_broadcast); -#endif - // inc6 (kEpFlat): write this token's metadata (SourceMeta + scales + topk) straight - // into each destination GPU's pool META region at the token's final recv slot, so a - // local consumer can produce recv_* without the forwarder + ring 2-hop. Mirrors the - // ep_dst_pools capture loop (warp-uniform over (node, g)); fields are written - // warp-cooperatively. topk_idx is stored RAW; the consumer rebases to its own expert - // range. The per-token __threadfence_system() in the epilogue publishes these writes. - // NOTE: only active under kEpFlat; the inc5 2-hop path (kEpFlat unset) is unchanged. - if (kEpFlat) { - // 8 (SourceMeta) + num_scales*4 + num_topk*4 (idx) + num_topk*4 (weights) must fit. - EP_DEVICE_ASSERT(static_cast(sizeof(SourceMeta)) + static_cast(num_scales) * 4 + - static_cast(num_topk) * 8 <= - kEpFlatMetaBytes); -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) { - const int node = topk_ranks[j]; - SourceMeta sm_j = warpBroadcast(src_meta, j); - NvlPackT bools_j = warpBroadcast(is_token_in_rank_uint64, node); - const bool* bvals_j = reinterpret_cast(&bools_j); -#pragma unroll - for (int g = 0; g < NUM_MAX_NVL_PEERS; ++g) { - if (not bvals_j[g]) continue; - const int dg = node * NUM_MAX_NVL_PEERS + g; - const int recv_idx = __shfl_sync(0xffffffff, ep_my_idx[g], node); - uint8_t* mbase = reinterpret_cast(recv_pool_global_ptrs[dg]) + ep_meta_base + - static_cast(recv_idx) * kEpFlatMetaBytes; - if (lane_id == 0) st_na_global(reinterpret_cast(mbase), sm_j); - float* sdst = reinterpret_cast(mbase + sizeof(SourceMeta)); - for (int s = lane_id; s < num_scales; s += 32) - st_na_global(sdst + s, ld_nc_global(x_scales + token_idx * num_scales + s)); - if (lane_id < num_topk) { - int* idst = reinterpret_cast(mbase + sizeof(SourceMeta) + num_scales * sizeof(float)); - float* wdst = reinterpret_cast(idst + num_topk); - st_na_global(idst + lane_id, static_cast(ld_nc_global(topk_idx + token_idx * num_topk + lane_id))); - st_na_global(wdst + lane_id, ld_nc_global(topk_weights + token_idx * num_topk + lane_id)); - } - } - } - __syncwarp(); - } - } else { - auto st_broadcast = [=](const int key, const int4& value) { -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) - st_na_global(reinterpret_cast(dst_send_buffers[j]) + key, value); - }; - UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast); - } -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) - dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + ring_hidden_int4; - - // Copy source metadata into symmetric send buffer - if (lane_id < num_topk_ranks) st_na_global(reinterpret_cast(dst_send_buffers[lane_id]), src_meta); -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) - dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + 1; - -// Copy `x_scales` into symmetric send buffer -#pragma unroll - for (int i = lane_id; i < num_scales; i += 32) { - auto value = ld_nc_global(x_scales + token_idx * num_scales + i); -#pragma unroll - for (int j = 0; j < num_topk_ranks; ++j) st_na_global(reinterpret_cast(dst_send_buffers[j]) + i, value); - } -#pragma unroll - for (int i = 0; i < num_topk_ranks; ++i) - dst_send_buffers[i] = reinterpret_cast(dst_send_buffers[i]) + num_scales; - -// Copy `topk_idx` and `topk_weights` into symmetric send buffer -#pragma unroll - for (int i = lane_id; i < num_topk * num_topk_ranks; i += 32) { - auto rank_idx = i / num_topk, copy_idx = i % num_topk; - auto idx_value = static_cast(ld_nc_global(topk_idx + token_idx * num_topk + copy_idx)); - auto weight_value = ld_nc_global(topk_weights + token_idx * num_topk + copy_idx); - st_na_global(reinterpret_cast(dst_send_buffers[rank_idx]) + copy_idx, idx_value); - st_na_global(reinterpret_cast(dst_send_buffers[rank_idx]) + num_topk + copy_idx, weight_value); - } - } - - // Epilogue - // Acquire sequential lock (lock path only; atomic-route removed the lock) - while (!ep_use_atomic_route and lane_id == 0 and rdma_send_next_token_idx != token_idx); - __syncwarp(); - - // Update last token tail (epilogue). Dead under kEpFlat (coordinator returns - // before reading rdma_send_channel_tail); kept for the non-flat path. - if (last_rdma_tail_idx >= 0) atomicMax(const_cast(rdma_send_channel_tail + lane_id), last_rdma_tail_idx + 1); - - // Release sequential lock (lock path only) - if (!ep_use_atomic_route and lane_id == 0) rdma_send_next_token_idx += 1; - -#if EP_NCCLEP_TMA - // B-depth-1: single drain of all cross-token in-flight TMA S2G groups at warp end - // (replaces the per-token wait_group 0), then publish to the generic proxy before the - // metadata threadfence below so the post-barrier consumer sees complete hidden + meta. - if (kEpDirect) { - asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); - __syncwarp(); - asm volatile("fence.proxy.async.global;" ::: "memory"); - } -#endif - // inc5: flush the direct cross-node hidden writes to the fabric so the - // destination ranks observe a complete recv_x after the dispatch barrier. - if (kEpDirect) __threadfence_system(); - } else if (warp_role == WarpRole::kRDMASenderCoordinator) { - // NOTES: in case of splitting the issued put at the end of the buffer - EP_DEVICE_ASSERT(num_max_rdma_chunked_recv_tokens % num_max_rdma_chunked_send_tokens == 0); - - // Synchronize shared memory - sync_rdma_sender_smem(); - // inc6 (kEpFlat): the all-sender path delivers hidden + metadata directly to - // the destination pools, so there is no RDMA ring to coordinate. The barrier - // above keeps the sender warps' bar.sync 0 balanced (7 warps); after it this - // coordinator warp exits. - if (kEpFlat) return; - if (lane_id == 0 && channel_id == 0 && rank == 0) { - (void)0; - } - - // Get number of tokens to send for each RDMA rank - int num_tokens_to_send = 0; - if (lane_id < kNumRDMARanks) { - num_tokens_to_send = rdma_channel_prefix_matrix[lane_id * num_channels + channel_id]; - if (channel_id > 0) num_tokens_to_send -= rdma_channel_prefix_matrix[lane_id * num_channels + channel_id - 1]; - } - - // Iterate all RDMA ranks - int last_issued_tail = 0; - while (__any_sync(0xffffffff, num_tokens_to_send > 0)) { -#pragma unroll - for (int i = 0; i < kNumRDMARanks; ++i, __syncwarp()) { - // To mitigate incast congestion, shuffle the starting index of target rank for different ranks and channels - const int dst_rdma_rank = (i + channel_id + rdma_rank) % kNumRDMARanks; - - // ----------------------------------------------------------------- - // Phase 4 restructure: owner lane (lane_id == dst_rdma_rank) decides - // whether to issue, then broadcasts num_tokens_to_issue and the - // tail-base for slot indexing to all 32 lanes via shfl. All lanes - // then participate in the cooperative cross-node fabric-IPC copy - // (when peer_rdma_bases != nullptr). This replaces the legacy - // single-lane handle.put(data) path which is broken on Azure CX-7. - // ----------------------------------------------------------------- - const bool owner = (lane_id == dst_rdma_rank); - int my_num_tokens_to_issue = 0; - int my_issue_tail = last_issued_tail; - if (owner && num_tokens_to_send > 0) { - auto processed_tail = ld_acquire_cta(const_cast(rdma_send_channel_tail + dst_rdma_rank)); - auto num_tokens_processed = processed_tail - last_issued_tail; - if (num_tokens_processed == num_tokens_to_send || num_tokens_processed >= num_max_rdma_chunked_send_tokens) { - int n = min(num_tokens_processed, num_max_rdma_chunked_send_tokens); - EP_DEVICE_ASSERT(n >= 0 && n <= num_tokens_to_send); - my_num_tokens_to_issue = n; - } - } - - const int n_issue = __shfl_sync(0xffffffff, my_num_tokens_to_issue, dst_rdma_rank); - const int issue_tail = __shfl_sync(0xffffffff, my_issue_tail, dst_rdma_rank); - if (n_issue == 0) continue; - - if (nvls_tail_mc != nullptr) { - // Phase 3+4: NVLS counter fast path. For cross-node, payload - // delivery uses warp-cooperative direct fabric-IPC writes when - // peer_rdma_bases is available; otherwise falls back to single- - // lane handle.put + flush. - if (dst_rdma_rank != rdma_rank) { - const auto dst_slot_idx = issue_tail % num_max_rdma_chunked_recv_tokens; - const size_t num_bytes_per_msg = (size_t)num_bytes_per_rdma_token * (size_t)n_issue; - const auto dst_offset = rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_recv_offset; - const auto src_offset = dst_rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_send_offset; - if (peer_rdma_bases != nullptr) { - // Phase 4: warp-cooperative int4 stores via NVL72 fabric VA. - // Uses NVLS counter (nvls_ctr_add) below to publish the - // writes to the receiver — `multimem.red.RELAXED` + a - // preceding `__threadfence_system()` is the validated - // ordering pair (release semantics on multimem.red triggers - // unspecified launch failure on Azure GB200). - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const int4* src_p = - reinterpret_cast(reinterpret_cast(rdma_buffer_ptr_base) + src_offset); - int4* dst_p = - reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + dst_offset); - const int n_int4 = (int)(num_bytes_per_msg / sizeof(int4)); - // Unrolled 8x to give the LSU pipeline more outstanding stores - // per lane. Each lane handles k, k+32, ..., k+224 per iter - // (stride 256) before looping. Tail handled with stride-32 loop. - const int stride8 = 8 * 32; - int k = lane_id; - const int n_full = (n_int4 / stride8) * stride8; - for (; k < n_full; k += stride8) { - int4 v0 = src_p[k]; - int4 v1 = src_p[k + 32]; - int4 v2 = src_p[k + 64]; - int4 v3 = src_p[k + 96]; - int4 v4 = src_p[k + 128]; - int4 v5 = src_p[k + 160]; - int4 v6 = src_p[k + 192]; - int4 v7 = src_p[k + 224]; - dst_p[k] = v0; - dst_p[k + 32] = v1; - dst_p[k + 64] = v2; - dst_p[k + 96] = v3; - dst_p[k + 128] = v4; - dst_p[k + 160] = v5; - dst_p[k + 192] = v6; - dst_p[k + 224] = v7; - } - for (; k < n_int4; k += 32) { - dst_p[k] = src_p[k]; - } - __syncwarp(); - __threadfence_system(); - } else if (owner) { - const auto port_channel_idx = - kLowLatencyMode ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - handle.put(dst_offset, src_offset, num_bytes_per_msg); - handle.flush(); - } - } - // Owner advances tail counter for this peer. - // Phase 4 fix: cross-node tail goes through direct fabric-VA - // store on peer's rdma_channel_tail slot. Self-loop tail goes - // through plain local atomicAdd on rdma_channel_tail.buffer(rdma_rank) - // — NVLS multicast is WRONG for self-loop because it fans out - // to all bound NVL peers' buffers (4 NVL ranks × n_issue ⇒ 4x - // over-count on each consumer's read). - if (owner) { - if (peer_rdma_bases != nullptr && dst_rdma_rank != rdma_rank) { - const int dst_rank_global = dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank; - const uintptr_t my_tail_off = reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - uint64_t* peer_tail = reinterpret_cast( - reinterpret_cast(peer_rdma_bases[dst_rank_global]) + my_tail_off); - const uint64_t new_tail = (uint64_t)issue_tail + (uint64_t)n_issue; - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(peer_tail), "l"(new_tail) : "memory"); - } else { - // Self-loop: plain release atomic on local slot (no multicast). - mscclpp::atomicFetchAdd(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)), - (uint64_t)n_issue, mscclpp::memoryOrderRelease); - } - } - } else if (owner) { - // Legacy non-NVLS path (single-lane). - if (dst_rdma_rank == rdma_rank) { - // Update tails - mscclpp::atomicFetchAdd(reinterpret_cast(rdma_channel_tail.buffer(rdma_rank)), (uint64_t)n_issue, - mscclpp::memoryOrderRelease); - } else { - const auto dst_slot_idx = issue_tail % num_max_rdma_chunked_recv_tokens; - const size_t num_bytes_per_msg = (size_t)num_bytes_per_rdma_token * (size_t)n_issue; - const auto dst_offset = rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_recv_offset; - const auto src_offset = dst_rdma_rank * (num_max_rdma_chunked_recv_tokens * num_bytes_per_rdma_token) + - dst_slot_idx * num_bytes_per_rdma_token + data_send_offset; - const auto port_channel_idx = kLowLatencyMode - ? (channel_id * kNumRDMARanks + dst_rdma_rank) - : (channel_id * num_ranks + dst_rdma_rank * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - handle.put(dst_offset, src_offset, num_bytes_per_msg); - - // HW atomicAdd is broken on Azure CX-7 RoCE (IBV_ATOMIC_NONE). - // Write the new absolute tail to a local scratch slot, then RDMA - // WRITE it to the peer's tail recv slot. - const uint64_t new_tail = (uint64_t)issue_tail + (uint64_t)n_issue; - *rdma_channel_tail_send_src.buffer(dst_rdma_rank) = new_tail; - __threadfence_system(); - const auto src_off_tail = reinterpret_cast(rdma_channel_tail_send_src.buffer(dst_rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - handle.put(rdma_rank * sizeof(uint64_t) + tail_send_offset, src_off_tail, sizeof(uint64_t)); - } - } - if (owner) { - last_issued_tail += n_issue; - num_tokens_to_send -= n_issue; - } - } - } - // Phase 4 diag: report final issued tail per (channel, dst_rdma_rank). - if (lane_id < kNumRDMARanks && rank == 0) { - (void)0; - } - } else if (warp_role == WarpRole::kRDMAAndNVLForwarder) { - // RDMA consumers and NVL producers - const auto dst_nvl_rank = target_rank; - const auto dst_rank = rdma_rank * NUM_MAX_NVL_PEERS + dst_nvl_rank; - const auto dst_rank_expert_begin = dst_rank * (num_experts / num_ranks); - const auto dst_rank_expert_end = dst_rank_expert_begin + (num_experts / num_ranks); - - // Wait counters to arrive - int num_tokens_to_recv_from_rdma = 0, src_rdma_channel_prefix = 0; - int my_start_sum = 0; // per-lane (src_rdma) channel-local start; used by same-GPU direct-write - EP_DEVICE_ASSERT(kNumRDMARanks <= 32); - auto start_time = clock64(); - if (lane_id < kNumRDMARanks) { - while (true) { - auto meta_0 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + dst_nvl_rank); - auto meta_1 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS + dst_nvl_rank); - auto meta_2 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2); - auto meta_3 = ld_volatile_global(rdma_channel_meta.recv_buffer(lane_id) + NUM_MAX_NVL_PEERS * 2 + 1); - if (meta_0 < 0 and meta_1 < 0 and meta_2 < 0 and meta_3 < 0) { - // Notify NVL ranks - int start_sum = -meta_0 - 1, end_sum = -meta_1 - 1; - EP_DEVICE_ASSERT(start_sum >= 0 and end_sum >= 0 and end_sum >= start_sum); - my_start_sum = start_sum; - st_relaxed_sys_global(nvl_channel_prefix_start.buffer() + lane_id, -start_sum - 1); - st_relaxed_sys_global(nvl_channel_prefix_end.buffer() + lane_id, -end_sum - 1); - - // Save RDMA channel received token count - src_rdma_channel_prefix = -meta_2 - 1; - auto src_rdma_channel_prefix_1 = -meta_3 - 1; - num_tokens_to_recv_from_rdma = src_rdma_channel_prefix_1 - src_rdma_channel_prefix; - if (not kCachedMode) - recv_rdma_channel_prefix_matrix[lane_id * num_channels + channel_id] = src_rdma_channel_prefix_1; - src_rdma_channel_prefix += lane_id == 0 ? 0 : recv_rdma_rank_prefix_sum[lane_id - 1]; - EP_DEVICE_ASSERT(num_tokens_to_recv_from_rdma >= 0); - // Phase 4 diag: report received expected token count per (channel, src). - if (rank == 4) { - (void)0; - } - break; - } - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch forwarder timeout (RDMA meta), channel: %d, RDMA: %d, nvl: %d, src RDMA lane: %d, dst " - "NVL: %d, meta: %d, %d, %d, %d\n", - channel_id, rdma_rank, nvl_rank, lane_id, dst_nvl_rank, meta_0, meta_1, meta_2, meta_3); - trap(); - } - } - } - __syncwarp(); - - // Shift cached head - send_nvl_head += src_rdma_channel_prefix * NUM_MAX_NVL_PEERS + dst_nvl_rank; - - // Wait shared memory to be cleaned - sync_forwarder_smem(); - - // ------------------------------------------------------------------- - // Increment 1 (NCCL-EP port): same-GPU fused direct-write. When this - // forwarder's destination NVL rank IS the local GPU (dst_nvl_rank == - // nvl_rank), the token arrived on the GPU it is destined for (no NVLink - // transpose). Instead of staging through nvl_channel_x and letting a - // separate receiver warp copy it to recv_x (the dest-side HBM double- - // bounce), write it straight to the final recv_x here; the matching - // receiver warp (src_nvl_rank == nvl_rank) is skipped. Final indices are - // computed identically to the receiver: recv_gbl_rank_prefix_sum base + - // channel-local start_sum + per-source running count (disjoint per - // src_global => no collision). combine's send_nvl_head breadcrumb is left - // unchanged, so combine for same-GPU tokens needs a paired change (gated - // for dispatch-only measurement). - const bool fuse_direct = (dst_nvl_rank == nvl_rank); - // Increment 3/4: cross-GPU forwarder direct-write. When the recv-output pool - // is active (recv_pool_ptrs != nullptr) and this warp's destination is a DIFFERENT - // NVL rank, write hidden straight to the destination peer's recv_x pool at - // the final index (peer-read prefix), removing the receiver's hidden drain. - const bool xnode_direct = (recv_pool_ptrs != nullptr) and (dst_nvl_rank != nvl_rank); - const int64_t ep_pool_header_bytes = ((static_cast(num_ranks) * 4 + 127) / 128) * 128; - int direct_recv_idx = 0; - if (lane_id < kNumRDMARanks and (fuse_direct or xnode_direct)) { - const int src_global = lane_id * NUM_MAX_NVL_PEERS + nvl_rank; - int gbl_base; - if (fuse_direct) { - gbl_base = (src_global > 0) ? recv_gbl_rank_prefix_sum[src_global - 1] : 0; - } else { - // Peer-read the destination GPU's recv_gbl_rank_prefix_sum from its pool - // header (published by the host before launch; visible by the time the - // forwarder drains, since all intra-node peers are co-launched). - const int* dst_prefix = reinterpret_cast(recv_pool_ptrs[dst_nvl_rank]); - gbl_base = (src_global > 0) ? static_cast(ld_volatile_global(dst_prefix + (src_global - 1))) : 0; - } - direct_recv_idx = gbl_base + my_start_sum; - if (fuse_direct and not kCachedMode) - recv_gbl_channel_prefix_matrix[src_global * num_channels + channel_id] = direct_recv_idx; - } - - // Forward tokens from RDMA buffer - // NOTES: always start from the local rank - int src_rdma_rank = sm_id % kNumRDMARanks; - int cached_rdma_channel_head = 0, cached_rdma_channel_tail = 0; - int cached_nvl_channel_head = 0, cached_nvl_channel_tail = 0, rdma_nvl_token_idx = 0; - while (__any_sync(0xffffffff, num_tokens_to_recv_from_rdma > 0)) { - // Check destination queue emptiness, or wait a buffer to be released - start_time = clock64(); - while (lane_id == 0) { - int num_used_slots = cached_nvl_channel_tail - cached_nvl_channel_head; - if (num_max_nvl_chunked_recv_tokens - num_used_slots >= num_max_nvl_chunked_send_tokens) break; - cached_nvl_channel_head = ld_volatile_global(nvl_channel_head.buffer()); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch forwarder timeout (NVL check), channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, head: %d, " - "tail: %d\n", - channel_id, rdma_rank, nvl_rank, dst_nvl_rank, ld_volatile_global(nvl_channel_head.buffer()), - cached_nvl_channel_tail); - trap(); - } - } - __syncwarp(); - - // Find next source RDMA rank (round-robin) - start_time = clock64(); - while (true) { - src_rdma_rank = (src_rdma_rank + 1) % kNumRDMARanks; - if (__shfl_sync(0xffffffff, num_tokens_to_recv_from_rdma, src_rdma_rank) > 0) { - if (lane_id == src_rdma_rank and cached_rdma_channel_head == cached_rdma_channel_tail) { - // Phase 4 fix: cross-node tail is now delivered via direct - // fabric-VA store to local rdma_channel_tail.buffer(src_rdma_rank), - // so prefer the legacy ld_acquire read which sees those stores. - // NVLS counter only used for self path (single rank). - // Phase 4 fix: cross-node tail comes via direct fabric-VA store - // (sender writes peer's rdma_channel_tail slot). Self-loop tail - // is a plain local atomic. Both end up in rdma_channel_tail — - // single read path via ld_acquire_sys_global covers them. - cached_rdma_channel_tail = static_cast(ld_acquire_sys_global(rdma_channel_tail.buffer(src_rdma_rank))); - } - if (__shfl_sync(0xffffffff, cached_rdma_channel_tail > cached_rdma_channel_head, src_rdma_rank)) break; - } - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES and lane_id < kNumRDMARanks) { - printf( - "DeepEP dispatch forwarder timeout (RDMA check), channel: %d, RDMA: %d, nvl: %d, dst NVL: %d, src RDMA " - "lane: %d, head: %d, tail: %d, expected: %d\n", - channel_id, rdma_rank, nvl_rank, dst_nvl_rank, lane_id, cached_rdma_channel_head, - cached_rdma_channel_tail, num_tokens_to_recv_from_rdma); - trap(); - } - } - auto src_rdma_head = __shfl_sync(0xffffffff, cached_rdma_channel_head, src_rdma_rank); - auto src_rdma_tail = __shfl_sync(0xffffffff, cached_rdma_channel_tail, src_rdma_rank); - - if (rank == 4 && lane_id == 0 && dst_nvl_rank == 0 && channel_id == 0) { - (void)0; - } - - // Iterate over every token from the RDMA buffer - for (int i = src_rdma_head, num_tokens_sent = 0; i < src_rdma_tail; ++i) { - auto rdma_slot_idx = i % num_max_rdma_chunked_recv_tokens; - void* shifted = rdma_channel_data.recv_buffer(src_rdma_rank) + rdma_slot_idx * num_bytes_per_rdma_token; - auto src_meta = - ld_nc_global(reinterpret_cast(reinterpret_cast(shifted) + ring_hidden_bytes)); - lane_id == src_rdma_rank ? (num_tokens_to_recv_from_rdma -= 1) : 0; - bool is_in_dst_nvl_rank = src_meta.is_token_in_nvl_rank(dst_nvl_rank); - if (lane_id == src_rdma_rank) { - auto cached_head = is_in_dst_nvl_rank ? rdma_nvl_token_idx : -1; - rdma_nvl_token_idx += is_in_dst_nvl_rank; - if (not kCachedMode) send_nvl_head[i * NUM_MAX_NVL_PEERS] = cached_head; - } - if (not is_in_dst_nvl_rank) continue; - - // Increment 1: same-GPU fused direct-write to final recv_x (skips - // nvl_channel_x staging + the separate receiver copy). - if (fuse_direct) { - int recv_token_idx = __shfl_sync(0xffffffff, direct_recv_idx, src_rdma_rank); - if (lane_id == src_rdma_rank) direct_recv_idx += 1; - - // Hidden (inc5: skipped when kEpDirect — the sender wrote it directly - // to the dest pool; here we only advance past the ring hidden region). - if (not kEpDirect) { - UNROLLED_WARP_COPY(28, lane_id, hidden_int4, recv_x + recv_token_idx * hidden_int4, - reinterpret_cast(shifted), ld_nc_global, st_na_global); - } - shifted = reinterpret_cast(shifted) + ring_hidden_int4; - - // Source meta - if (lane_id == 0 and not kCachedMode) st_na_global(recv_src_meta + recv_token_idx, src_meta); - shifted = reinterpret_cast(shifted) + 1; - - // x_scales - UNROLLED_WARP_COPY(1, lane_id, num_scales, recv_x_scales + recv_token_idx * num_scales, - reinterpret_cast(shifted), ld_nc_global, st_na_global); - shifted = reinterpret_cast(shifted) + num_scales; - - // topk_idx / topk_weights (rebased to dst-rank-local expert range) - if (lane_id < num_topk) { - auto idx_value = ld_nc_global(reinterpret_cast(shifted) + lane_id); - shifted = reinterpret_cast(shifted) + num_topk; - auto weight_value = ld_nc_global(reinterpret_cast(shifted) + lane_id); - idx_value = (idx_value >= dst_rank_expert_begin and idx_value < dst_rank_expert_end) - ? idx_value - dst_rank_expert_begin - : -1; - st_na_global(recv_topk_idx + recv_token_idx * num_topk + lane_id, static_cast(idx_value)); - weight_value = idx_value >= 0 ? weight_value : 0.0f; - st_na_global(recv_topk_weights + recv_token_idx * num_topk + lane_id, weight_value); - } - continue; - } - - // Get an empty slot (still used for scales/topk/meta staging that the - // receiver drains; under xnode_direct the hidden goes straight to the pool) - int dst_slot_idx = (cached_nvl_channel_tail++) % num_max_nvl_chunked_recv_tokens; - - // Hidden copy: inc5 (kEpDirect) skips it entirely (sender wrote the dest - // pool directly); else increment 3 cross-GPU direct-write to the dest's - // recv_x pool at the final index; else legacy nvl_channel staging. - if (kEpDirect) { - // inc5: hidden already in the dest pool (sender direct-write); skip. - } else if (xnode_direct) { - int recv_token_idx = __shfl_sync(0xffffffff, direct_recv_idx, src_rdma_rank); - if (lane_id == src_rdma_rank) direct_recv_idx += 1; -#if EP_NCCLEP_TMA_LOCALDST - // DIAGNOSTIC: target the LOCAL (non-IPC-imported) VMM pool VA to isolate - // whether the TMA launch failure is specific to cross-process IPC peer - // targets. Corrupts results (expected); only the launch outcome matters. - uint8_t* dst_b = reinterpret_cast(recv_pool_ptrs[nvl_rank]) + ep_pool_header_bytes + - static_cast(recv_token_idx) * hidden_bytes; -#else - uint8_t* dst_b = reinterpret_cast(recv_pool_ptrs[dst_nvl_rank]) + ep_pool_header_bytes + - static_cast(recv_token_idx) * hidden_bytes; -#endif -#if EP_NCCLEP_TMA - // Stream HBM(shifted) -> SMEM -> peer VMM VA via cp.async.bulk (TMA). - // Proven crash-free against imported cuMem FABRIC/POSIX-FD peer VAs - // (would crash for cudaIpc-mapped buffer_ptrs). lane 0 drives the warp's - // tile; chunked to bound SMEM. fence.proxy orders async-proxy writes. - const uint8_t* src_b = reinterpret_cast(shifted); - if (lane_id == 0) { - const uint32_t smem_tile = static_cast(__cvta_generic_to_shared(ep_tma_tile[warp_id])); - const uint32_t smem_mbar = static_cast(__cvta_generic_to_shared(&ep_tma_mbar[warp_id])); - size_t off = 0; - const size_t total = static_cast(hidden_bytes); - while (off < total) { - const size_t remain = total - off; - uint32_t chunk = - static_cast(remain < (size_t)kEpTmaTileBytes ? remain : (size_t)kEpTmaTileBytes); - chunk = (chunk + 15u) & ~15u; - // Mirror the validated probe sequence exactly: init -> fence -> - // G2S bulk (mbarrier complete_tx) -> arrive.expect_tx -> wait -> - // fence -> S2G bulk -> commit -> wait. - asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_mbar)); - fenceProxyAsyncSharedCta(); - asm volatile( - "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " - "[%0], [%1], %2, [%3];" ::"r"(smem_tile), - "l"(src_b + off), "r"(chunk), "r"(smem_mbar) - : "memory"); - uint64_t st; - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 %0, [%1], %2;" - : "=l"(st) - : "r"(smem_mbar), "r"(chunk)); - uint32_t done = 0; - while (!done) { - asm volatile( - "{ .reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], 0;" - " selp.u32 %0, 1, 0, p; }" - : "=r"(done) - : "r"(smem_mbar)); - } - fenceProxyAsyncSharedCta(); - asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst_b + off), - "r"(smem_tile), "r"(chunk) - : "memory"); - asm volatile("cp.async.bulk.commit_group;"); - asm volatile("cp.async.bulk.wait_group 0;"); - off += chunk; - } - asm volatile("fence.proxy.async.global;" ::: "memory"); - } - __syncwarp(); -#else - int4* dst_pool_x = reinterpret_cast(dst_b); - UNROLLED_WARP_COPY(28, lane_id, hidden_int4, dst_pool_x, reinterpret_cast(shifted), ld_nc_global, - st_na_global); -#endif - } else { - UNROLLED_WARP_COPY(28, lane_id, hidden_int4, nvl_channel_x.buffer() + dst_slot_idx * hidden_int4, - reinterpret_cast(shifted), ld_nc_global, st_na_global); - } - shifted = reinterpret_cast(shifted) + ring_hidden_int4; - - // Copy source meta - if (lane_id == 0) st_na_global(nvl_channel_src_meta.buffer() + dst_slot_idx, src_meta); - shifted = reinterpret_cast(shifted) + 1; - - // Copy `x_scales` - UNROLLED_WARP_COPY(1, lane_id, num_scales, nvl_channel_x_scales.buffer() + dst_slot_idx * num_scales, - reinterpret_cast(shifted), ld_nc_global, st_na_global); - shifted = reinterpret_cast(shifted) + num_scales; - - // Copy `topk_idx` and `topk_weights` - // NOTES: do not use `shifted` after this `if`, because only several lanes are shifted - if (lane_id < num_topk) { - // Read - auto idx_value = ld_nc_global(reinterpret_cast(shifted) + lane_id); - shifted = reinterpret_cast(shifted) + num_topk; - auto weight_value = ld_nc_global(reinterpret_cast(shifted) + lane_id); - - // Transform and write - idx_value = (idx_value >= dst_rank_expert_begin and idx_value < dst_rank_expert_end) - ? idx_value - dst_rank_expert_begin - : -1; - st_na_global(nvl_channel_topk_idx.buffer() + dst_slot_idx * num_topk + lane_id, idx_value); - weight_value = idx_value >= 0 ? weight_value : 0.0f; - st_na_global(nvl_channel_topk_weights.buffer() + dst_slot_idx * num_topk + lane_id, weight_value); - } - - // In case of insufficient NVL buffers, early stopping - if ((++num_tokens_sent) == num_max_nvl_chunked_send_tokens) src_rdma_tail = i + 1; - } - - // Sync head index - if (lane_id == src_rdma_rank) - forward_channel_head[dst_nvl_rank][src_rdma_rank] = (cached_rdma_channel_head = src_rdma_tail); - - // Move tail index - __syncwarp(); - if (lane_id == 0) st_release_sys_global(nvl_channel_tail.buffer(), cached_nvl_channel_tail); - } - - // Retired - __syncwarp(); - if (lane_id == 0) { - forward_channel_retired[dst_nvl_rank] = true; - if (channel_id == 0) { - (void)0; - } - } - } else if (warp_role == WarpRole::kForwarderCoordinator) { - // Extra warps for forwarder coordinator should exit directly - if (target_rank > 0) return; - - // Forward warp coordinator - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - - // Clean shared memory - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS <= 32, "Invalid number of NVL peers"); -#pragma unroll - for (int i = lane_id; i < kNumRDMARanks * NUM_MAX_NVL_PEERS; i += 32) - forward_channel_head[i % NUM_MAX_NVL_PEERS][i / NUM_MAX_NVL_PEERS] = 0; - if (lane_id < NUM_MAX_NVL_PEERS) forward_channel_retired[lane_id] = false; - sync_forwarder_smem(); - - int last_head = 0, target_rdma = lane_id < kNumRDMARanks ? lane_id : 0; - while (true) { - // Find minimum head - int min_head = std::numeric_limits::max(); -#pragma unroll - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) - if (not forward_channel_retired[i]) min_head = min(min_head, forward_channel_head[i][target_rdma]); - if (__all_sync(0xffffffff, min_head == std::numeric_limits::max())) break; - - // Update remote head - // Phase 4 perf: lower the lazy-update threshold from chunk_send → chunk_send/4 - // (or 1 if retired) so the sender's receive-buffer-space window - // advances more frequently. The original threshold caused partial - // last chunks to never trigger head feedback, which deadlocked at - // larger chunk_send values where 4096 tokens / 10 channels / 2 peers - // ≈ 205 tokens ⇒ 3 full chunks + 1 partial. - const bool any_retired = forward_channel_retired[0] || forward_channel_retired[1] || forward_channel_retired[2] || - forward_channel_retired[3] || forward_channel_retired[4] || forward_channel_retired[5] || - forward_channel_retired[6] || forward_channel_retired[7]; - const int head_update_threshold = any_retired ? 1 : max(1, num_max_rdma_chunked_send_tokens / 4); - if (min_head != std::numeric_limits::max() and min_head >= last_head + head_update_threshold and - lane_id < kNumRDMARanks) { - if (peer_rdma_bases != nullptr && lane_id != rdma_rank) { - // Phase 4 fix: cross-node head feedback via direct fabric-VA - // store on peer's rdma_channel_head slot (single writer per - // (channel, my_rdma_rank) on peer's side: producer is me, slot - // is `peer.rdma_channel_head.buffer(my_rdma_rank)`). Bypasses - // both broken port_channel.put and the unreliable NVLS counter. - const int dst_rank_global = lane_id * NUM_MAX_NVL_PEERS + nvl_rank; - const uintptr_t my_head_off = reinterpret_cast(rdma_channel_head.buffer(rdma_rank)) - - reinterpret_cast(rdma_buffer_ptr_base); - uint64_t* peer_head = - reinterpret_cast(reinterpret_cast(peer_rdma_bases[dst_rank_global]) + my_head_off); - asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(peer_head), "l"((uint64_t)min_head) : "memory"); - } else if (lane_id == rdma_rank) { - // Self-loop: plain release atomic on local slot. Cannot use NVLS - // multimem here \u2014 it fans out to all NVL peers' local buffers - // and over-counts (4 NVL ranks \u00d7 add \u21d2 4x increment). - mscclpp::atomicFetchAdd(static_cast(rdma_channel_head.buffer(rdma_rank)), - (uint64_t)(min_head - last_head), mscclpp::memoryOrderRelease); - } else { - auto dst_offset = rdma_rank * sizeof(uint64_t) + head_send_offset; - auto port_channel_idx = kLowLatencyMode ? (channel_id * kNumRDMARanks + lane_id) - : (channel_id * num_ranks + lane_id * NUM_MAX_NVL_PEERS + nvl_rank); - auto& handle = port_channel_handles[port_channel_idx]; - // Absolute-value RDMA WRITE replaces broken HW atomicAdd (see note above). - *rdma_channel_head_send_src.buffer(lane_id) = (uint64_t)min_head; - __threadfence_system(); - const auto src_off_head = reinterpret_cast(rdma_channel_head_send_src.buffer(lane_id)) - - reinterpret_cast(rdma_buffer_ptr_base); - handle.put(dst_offset, src_off_head, sizeof(uint64_t)); - } - last_head = min_head; - } - - // Nanosleep and let other warps work - __nanosleep(NUM_WAIT_NANOSECONDS); - } - } else { - // NVL consumers - // inc6 (kEpFlat): the all-sender path has no forwarder feeding the NVL receive - // buffers, so these consumer warps would spin to the timeout trap. The flat - // metadata drain (flat_meta_drain) produces recv_* from the pool instead. - if (kEpFlat) return; - // Retrieve rank offset from barrier results (each lane's register stores an RDMA rank) - int src_nvl_rank = target_rank, total_offset = 0; - // Increment 1: same-GPU tokens (src_nvl_rank == nvl_rank) are written - // directly to recv_x by the fused forwarder above; skip this receiver warp. - if (src_nvl_rank == nvl_rank) return; - EP_STATIC_ASSERT(kNumRDMARanks <= 32, "Invalid number of RDMA peers"); - if (lane_id < kNumRDMARanks and lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank > 0) - total_offset = recv_gbl_rank_prefix_sum[lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank - 1]; - - // Receive channel offsets - int start_offset = 0, end_offset = 0, num_tokens_to_recv; - auto start_time = clock64(); - while (lane_id < kNumRDMARanks) { - start_offset = ld_volatile_global(nvl_channel_prefix_start.buffer() + lane_id); - end_offset = ld_volatile_global(nvl_channel_prefix_end.buffer() + lane_id); - if (start_offset < 0 and end_offset < 0) { - start_offset = -start_offset - 1, end_offset = -end_offset - 1; - total_offset += start_offset; - break; - } - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src RDMA: %d, src nvl: %d, start: " - "%d, end: %d\n", - channel_id, rdma_rank, nvl_rank, lane_id, src_nvl_rank, start_offset, end_offset); - trap(); - } - } - num_tokens_to_recv = warp_reduce_sum(end_offset - start_offset); - - // Save for combine usage - if (lane_id < kNumRDMARanks and not kCachedMode) - recv_gbl_channel_prefix_matrix[(lane_id * NUM_MAX_NVL_PEERS + src_nvl_rank) * num_channels + channel_id] = - total_offset; - __syncwarp(); - - int cached_channel_head_idx = 0, cached_channel_tail_idx = 0; - while (num_tokens_to_recv > 0) { - // Check channel status by lane 0 - start_time = clock64(); - while (lane_id == 0) { - // Ready to copy - if (cached_channel_head_idx != cached_channel_tail_idx) break; - cached_channel_tail_idx = ld_acquire_sys_global(nvl_channel_tail.buffer()); - - // Timeout check - if (clock64() - start_time > NUM_TIMEOUT_CYCLES) { - printf( - "DeepEP dispatch NVL receiver timeout, channel: %d, RDMA: %d, nvl: %d, src NVL: %d, head: %d, tail: %d\n", - channel_id, rdma_rank, nvl_rank, src_nvl_rank, cached_channel_head_idx, cached_channel_tail_idx); - trap(); - } - } - - // Sync queue tail - cached_channel_tail_idx = __shfl_sync(0xffffffff, cached_channel_tail_idx, 0); - - // Copy data - int num_recv_tokens = cached_channel_tail_idx - cached_channel_head_idx; - for (int chunk_idx = 0; chunk_idx < num_recv_tokens; ++chunk_idx, --num_tokens_to_recv) { - int token_idx_in_buffer = (cached_channel_head_idx++) % num_max_nvl_chunked_recv_tokens; - auto meta = ld_nc_global(nvl_channel_src_meta.buffer() + token_idx_in_buffer); - int64_t recv_token_idx = __shfl_sync(0xffffffff, total_offset, meta.src_rdma_rank); - (lane_id == meta.src_rdma_rank) ? (total_offset += 1) : 0; - -#if !EP_NCCLEP_DRAIN_NOOP - // Copy data (hidden) — skipped under increment-3 cross-GPU direct-write: - // the producing forwarder already wrote hidden straight to this GPU's - // recv_x pool at recv_token_idx. Scales/topk/meta still drain here. - if (recv_pool_ptrs == nullptr) { - UNROLLED_WARP_COPY(28, lane_id, hidden_int4, recv_x + recv_token_idx * hidden_int4, - nvl_channel_x.buffer() + token_idx_in_buffer * hidden_int4, ld_nc_global, st_na_global); - } - - // Copy source meta - if (lane_id == 0 and not kCachedMode) st_na_global(recv_src_meta + recv_token_idx, meta); - - // Copy scales - UNROLLED_WARP_COPY(1, lane_id, num_scales, recv_x_scales + recv_token_idx * num_scales, - nvl_channel_x_scales.buffer() + token_idx_in_buffer * num_scales, ld_nc_global, - st_na_global); - - // Copy `topk_idx` and `topk_weights` - if (lane_id < num_topk) { - auto recv_idx = recv_token_idx * num_topk + lane_id; - auto buffer_idx = token_idx_in_buffer * num_topk + lane_id; - st_na_global(recv_topk_idx + recv_idx, - static_cast(ld_nc_global(nvl_channel_topk_idx.buffer() + buffer_idx))); - st_na_global(recv_topk_weights + recv_idx, ld_nc_global(nvl_channel_topk_weights.buffer() + buffer_idx)); - } -#endif // !EP_NCCLEP_DRAIN_NOOP - } - - // Move queue - __syncwarp(); - if (lane_id == 0) st_relaxed_sys_global(nvl_channel_head.buffer(), cached_channel_head_idx); - } - } - if (thread_id == 0 && channel_id == 0) { - (void)0; - } -} - -#endif // EP_DISPATCH_NCCLEP -#endif // MSCCLPP_EP_INTERNODE_NCCLEP_CUH_ diff --git a/src/ext/ep/ht/layout.cu b/src/ext/ep/ht/layout.cu new file mode 100644 index 00000000..d88ccf58 --- /dev/null +++ b/src/ext/ep/ht/layout.cu @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP) +// branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. +// +// HT routing layout construction. + +#include "exception.cuh" +#include "launch.cuh" + +namespace mscclpp { +namespace ep { +namespace intranode { + +template +__global__ void __launch_bounds__(NumThreads, 1) + get_dispatch_layout(const int64_t* topk_idx, int* num_tokens_per_rank, int* num_tokens_per_expert, + bool* is_token_in_rank, int num_tokens, int num_topk, int num_ranks, int num_experts) { + const int block_id = static_cast(blockIdx.x); + const int thread_id = static_cast(threadIdx.x); + + union SharedStorage { + int perExpert[NumThreads][NumExpertsPerBlock]; + int perRank[NumThreads][NumRanksPerBlock]; + }; + __shared__ SharedStorage shared; + const int expert_begin = block_id * NumExpertsPerBlock; + const int expert_end = min(expert_begin + NumExpertsPerBlock, num_experts); + if (expert_begin < expert_end) { +#pragma unroll + for (int i = 0; i < NumExpertsPerBlock; ++i) shared.perExpert[thread_id][i] = 0; +#pragma unroll + for (int token = thread_id; token < num_tokens; token += NumThreads) { + const int64_t* token_topk = topk_idx + token * num_topk; +#pragma unroll + for (int i = 0; i < num_topk; ++i) { + const int expert = static_cast(token_topk[i]); + if (expert_begin <= expert and expert < expert_end) ++shared.perExpert[thread_id][expert - expert_begin]; + } + } + __syncthreads(); + + EP_STATIC_ASSERT(NumExpertsPerBlock <= NumThreads, "Too many experts per block"); + if (expert_begin + thread_id < expert_end) { + int sum = 0; +#pragma unroll + for (int i = 0; i < NumThreads; ++i) sum += shared.perExpert[i][thread_id]; + num_tokens_per_expert[expert_begin + thread_id] = sum; + } + return; + } + + const int rank_blocks = (num_experts + NumExpertsPerBlock - 1) / NumExpertsPerBlock; + const int rank_begin = (block_id - rank_blocks) * NumRanksPerBlock; + const int rank_end = min(rank_begin + NumRanksPerBlock, num_ranks); + if (rank_begin >= rank_end) return; + + const int experts_per_rank = num_experts / num_ranks; + const int rank_expert_begin = rank_begin * experts_per_rank; + const int rank_expert_end = rank_end * experts_per_rank; +#pragma unroll + for (int i = 0; i < NumRanksPerBlock; ++i) shared.perRank[thread_id][i] = 0; +#pragma unroll + for (int token = thread_id; token < num_tokens; token += NumThreads) { + const int64_t* token_topk = topk_idx + token * num_topk; + int is_in_rank[NumRanksPerBlock] = {0}; +#pragma unroll + for (int i = 0; i < num_topk; ++i) { + const int expert = static_cast(token_topk[i]); + if (rank_expert_begin <= expert and expert < rank_expert_end) + ++is_in_rank[expert / experts_per_rank - rank_begin]; + } + + bool* token_in_rank = is_token_in_rank + token * num_ranks; +#pragma unroll + for (int i = 0; rank_begin + i < rank_end; ++i) { + token_in_rank[rank_begin + i] = is_in_rank[i] > 0; + shared.perRank[thread_id][i] += is_in_rank[i] > 0; + } + } + __syncthreads(); + + EP_STATIC_ASSERT(NumRanksPerBlock <= NumThreads, "Too many ranks per block"); + if (rank_begin + thread_id < rank_end) { + int sum = 0; +#pragma unroll + for (int i = 0; i < NumThreads; ++i) sum += shared.perRank[i][thread_id]; + num_tokens_per_rank[rank_begin + thread_id] = sum; + } +} + +void get_dispatch_layout(const int64_t* topk_idx, int* num_tokens_per_rank, int* num_tokens_per_expert, + bool* is_token_in_rank, int num_tokens, int num_topk, int num_ranks, int num_experts, + cudaStream_t stream) { + constexpr int NumThreads = 256; + constexpr int NumExpertsPerBlock = 32; + constexpr int NumRanksPerBlock = 8; + const int num_blocks = (num_experts + NumExpertsPerBlock - 1) / NumExpertsPerBlock + + (num_ranks + NumRanksPerBlock - 1) / NumRanksPerBlock; + + SETUP_LAUNCH_CONFIG(num_blocks, NumThreads, stream); + LAUNCH_KERNEL(&cfg, (get_dispatch_layout), topk_idx, + num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank, num_tokens, num_topk, num_ranks, + num_experts); +} + +} // namespace intranode +} // namespace ep +} // namespace mscclpp diff --git a/src/ext/ep/ht/kernels/runtime.cu b/src/ext/ep/ht/runtime.cu similarity index 77% rename from src/ext/ep/ht/kernels/runtime.cu rename to src/ext/ep/ht/runtime.cu index 9d473ef8..36a9d433 100644 --- a/src/ext/ep/ht/kernels/runtime.cu +++ b/src/ext/ep/ht/runtime.cu @@ -4,10 +4,8 @@ // Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), // branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. // -// Intranode runtime helpers. Only the NVLink barrier launcher is ported here -// (see DeepEP `csrc/kernels/runtime.cu::intranode::barrier`). The -// internode/NVSHMEM init helpers are deliberately omitted; the MSCCL++ port -// uses `mscclpp::Bootstrap`/`ProxyService` instead of NVSHMEM. +// HT runtime helpers. This file contains the device barrier launcher +// used to coordinate direct peer access to the symmetric HT buffers. #include "constants.cuh" #include "device_helpers.cuh" diff --git a/src/ext/ep/ht_runtime.cc b/src/ext/ep/ht_runtime.cc index 9a56d441..b427d402 100644 --- a/src/ext/ep/ht_runtime.cc +++ b/src/ext/ep/ht_runtime.cc @@ -3,1645 +3,351 @@ // // Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP) // branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. -// -// Torch-free high-throughput MoE dispatch/combine runtime. This is a -// mechanical de-torch of the former `src/ext/ep/ht/buffer.cc`: tensor params -// became raw device pointers + explicit size ints, output tensors became -// caller-provided pointers, `at::cuda::CUDAStream` became `cudaStream_t`, and -// the torch validation asserts / EventHandle async machinery were dropped. The -// CUDA kernels (`intranode::` / `internode::`), env knobs, `#ifdef -// EP_DISPATCH_NCCLEP` blocks, recv-pool / VMM / IPC logic and diagnostics are -// preserved verbatim. #include "ht_runtime.hpp" -#include -#include - #include -#include -#include #include -#include #include -#include #include #include -#include -#include #include #include -#include -#include #include "api.cuh" #include "constants.cuh" +#include "exception.cuh" namespace mscclpp { namespace ep { -// Upstream MSCCL++ now exposes `Connection::atomicAdd` and -// `PortChannelDeviceHandle::atomicAdd` natively (see commit "atomic add" -// on branch chhwang/new-atomic-add, merged into this tree). The stock -// `ProxyService` recognises `ChannelTrigger.type == 0` as an atomic-add -// request, so no subclass or private-member access is required anymore. -using EPProxyService = mscclpp::ProxyService; +MoEHighThroughputRuntime::MoEHighThroughputRuntime(mscclpp::Communicator& communicator, int64_t maxHiddenBytes, + const Config& config) + : rank_(communicator.bootstrap()->getRank()), + numRanks_(communicator.bootstrap()->getNranks()), + numNvlRanks_(std::min(numRanks_, communicator.bootstrap()->getNranksPerNode())), + numRanksPerIpcDomain_( + std::max(numNvlRanks_, std::min(numRanks_, communicator.bootstrap()->getNranksPerIpcDomain()))), + maxHiddenBytes_(maxHiddenBytes), + config_(config) { + EP_HOST_ASSERT(rank_ >= 0 and rank_ < numRanks_); + EP_HOST_ASSERT(numNvlRanks_ > 0); + EP_HOST_ASSERT(maxHiddenBytes_ > 0); -// Number of host-side proxy services (== proxy threads) the runtime creates. -// PortChannels are sharded across these so per-thread FIFO contention drops. -// A single proxy thread caps cross-node LL combine at ~2.8 ms/iter on H100+IB -// platforms; 8 threads bring it down to ~470 us (NIC-bound). On NVSwitch-only -// platforms (e.g. GB200 NVL72) host proxies don't post IB WRs, so 1 is plenty. -// -// Resolution order: -// 1. `MSCCLPP_EP_NUM_PROXIES` env var (clamped to >= 1) if set. -// 2. Auto-detect from current device's compute capability: -// - Hopper (sm_90, H100/H200) and earlier: 8 -// - Blackwell (sm_100+, B100/B200/GB200): 1 -// 3. Fallback: 8. -// -// Empirical sweep on 2x8 H100 + IB (16 ranks, t=128, h=7168, topk=8): -// N=1: D 1013 us / C 2801 us -// N=2: D 611 us / C 843 us -// N=4: D 479 us / C 568 us -// N=8: D 445 us / C 469 us <-- knee -// N=12: collapses (CPU oversubscription with 8 GPUs/node). -static int resolve_num_proxy_services(int num_ranks, int local_world_size) { - if (const char* env = std::getenv("MSCCLPP_EP_NUM_PROXIES")) { - int v = std::atoi(env); - return v > 0 ? v : 1; - } - int dev = 0; - if (cudaGetDevice(&dev) != cudaSuccess) return 8; - cudaDeviceProp prop{}; - if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess) return 8; - // sm_100+ = Blackwell (GB200 etc.). NVSwitch fabric makes one proxy - // sufficient intranode, but cross-node RDMA still funnels through that - // single host thread. When num_rdma_ranks > 1, shard the proxy across - // local GPUs so each GPU has its own RDMA-driving thread. (P2) - int lws = local_world_size > 0 ? local_world_size : 1; - int num_rdma_ranks_local = std::max(1, num_ranks / lws); - if (prop.major >= 10) return num_rdma_ranks_local > 1 ? lws : 1; - return 8; -} + if ((numRanks_ != 2 and numRanks_ != 4 and numRanks_ != 8) or numRanksPerIpcDomain_ < numRanks_) return; -// Cross-host cuMem fabric IPC capability. -// -// `isNvlsSupported()` returns true for any device with NVLink multicast, -// including H100. But NVLS by itself only works inside one host's NVLink -// island; cross-host sharing of cuMem allocations / NVLS multicast handles -// requires the device to be on an actual cross-host NVLink fabric (GB200 -// NVL72 with nvidia-imex on Azure today). H100 reports -// `CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED == 1` too but lacks -// the NVSwitch fabric to actually share fabric handles across hosts; the -// cross-host import then falls back to POSIX-FD, whose handle exchange -// routes through a unix-domain socket on the master host -- which -// worker-node ranks cannot reach (`connect() failed for unix socket to -// /tmp/mscclpp_bootstrap_*.sock`). That is the exact failure signature -// commit 3ab2e43b ("NVLS HT B2 phases 1-3") introduced on H100 / Azure -// CX-7 RoCE. -// -// To stay safe by default we require both: -// - device attr `CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED` -// - compute capability >= sm_100 (Blackwell+). -// -// Resolution order: -// 1. `MSCCLPP_EP_FABRIC_IPC` env var (`0`/`off`/`false`/`no` => off, -// `1`/`on`/`true`/`yes`/`force` => on, anything else => auto). When -// set, env value takes precedence over the device check. -// 2. Auto-detect via the two checks above. -static bool resolve_fabric_ipc_supported() { - if (const char* env = std::getenv("MSCCLPP_EP_FABRIC_IPC")) { - std::string v(env); - for (auto& c : v) c = std::tolower(static_cast(c)); - if (v == "0" || v == "off" || v == "false" || v == "no") return false; - if (v == "1" || v == "on" || v == "true" || v == "yes" || v == "force") return true; - } - int dev = 0; - if (cudaGetDevice(&dev) != cudaSuccess) return false; - CUdevice cuDev; - if (cuDeviceGet(&cuDev, dev) != CUDA_SUCCESS) return false; - int supported = 0; - if (cuDeviceGetAttribute(&supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cuDev) != CUDA_SUCCESS) { - return false; - } - if (!supported) return false; - cudaDeviceProp prop{}; - if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess) return false; - // Blackwell+ (sm_100, GB200 NVL72) is the only deployed cross-host - // NVLink fabric today. H100 (sm_90) advertises fabric-handle support - // but lacks the nvidia-imex / NVSwitch fabric to actually share them - // across hosts. - return prop.major >= 10; -} - -MoEHighThroughputRuntime::MoEHighThroughputRuntime(int rank, int numRanks, int64_t numNvlBytes, int64_t numRdmaBytes) - : num_nvl_bytes(numNvlBytes), - num_rdma_bytes(numRdmaBytes), - rank(rank), - num_ranks(numRanks), - bootstrap(std::make_shared(rank, numRanks)) { - // Communication stream + reusable cross-stream ordering event. Created first - // because the IPC/workspace setup below issues `cudaMemsetAsync` on it. This - // replaces the torch `at::cuda::getStreamFromPool(true)` member initializer. - CUDA_CHECK(cudaStreamCreateWithFlags(&comm_stream, cudaStreamNonBlocking)); - CUDA_CHECK(cudaEventCreateWithFlags(&comm_event, cudaEventDisableTiming)); - - // Resolve local_world_size early so the Blackwell proxy-sharding (P2) can use it. - int local_world_size_for_proxy = NUM_MAX_NVL_PEERS; - if (const char* env = std::getenv("MSCCLPP_EP_LOCAL_WORLD_SIZE")) { - int v = std::atoi(env); - if (v > 0 && v <= NUM_MAX_NVL_PEERS) local_world_size_for_proxy = v; - } - num_proxy_services = resolve_num_proxy_services(num_ranks, local_world_size_for_proxy); - proxy_services.reserve(num_proxy_services); - for (int i = 0; i < num_proxy_services; ++i) { - proxy_services.emplace_back(std::make_shared()); - } - if (rank == 0) { - printf("[mscclpp_ep] num_proxy_services=%d (set MSCCLPP_EP_NUM_PROXIES to override)\n", num_proxy_services); - fflush(stdout); - } - // Task fifo memory - int64_t fifo_bytes = sizeof(int) * NUM_MAX_FIFO_SLOTS; - int64_t buffer_ptr_bytes = sizeof(void*) * NUM_MAX_NVL_PEERS; - int64_t task_ptr_bytes = sizeof(int*) * NUM_MAX_NVL_PEERS; - - // Common checks - EP_HOST_ASSERT(num_nvl_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 and - (num_nvl_bytes <= std::numeric_limits::max() or num_rdma_bytes == 0)); - EP_HOST_ASSERT(num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 and - (low_latency_mode or num_rdma_bytes <= std::numeric_limits::max())); - EP_HOST_ASSERT(0 <= rank and rank < num_ranks and - (num_ranks <= NUM_MAX_NVL_PEERS * NUM_MAX_RDMA_PEERS or low_latency_mode)); - EP_HOST_ASSERT(num_ranks < NUM_MAX_NVL_PEERS or num_ranks % NUM_MAX_NVL_PEERS == 0); - if (num_rdma_bytes > 0) EP_HOST_ASSERT(num_ranks > NUM_MAX_NVL_PEERS or low_latency_mode); - - // Get ranks - CUDA_CHECK(cudaGetDevice(&device_id)); - // Allow overriding the local-world-size (number of GPUs per node) via the - // env var MSCCLPP_EP_LOCAL_WORLD_SIZE. By default the partitioning is - // pinned to NUM_MAX_NVL_PEERS=8, which mis-classifies all ranks as - // intra-node on hosts with fewer than 8 GPUs (e.g. GB200x4) and breaks - // cross-node LL via spurious cudaIpcOpenMemHandle on remote IPC handles. - int local_world_size = NUM_MAX_NVL_PEERS; - if (const char* env = std::getenv("MSCCLPP_EP_LOCAL_WORLD_SIZE")) { - int v = std::atoi(env); - if (v > 0 && v <= NUM_MAX_NVL_PEERS) local_world_size = v; - } - rdma_rank = rank / local_world_size, nvl_rank = rank % local_world_size; - num_rdma_ranks = std::max(1, num_ranks / local_world_size), num_nvl_ranks = std::min(num_ranks, local_world_size); - - // Get device info - cudaDeviceProp device_prop = {}; - CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_id)); - - if (num_nvl_bytes > 0) { - // Local IPC: alloc local memory and set local IPC handle - int64_t ep_nvl_alloc = num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes; -#ifdef EP_DISPATCH_NCCLEP - // Increment 4: allocate the peer-mapped recv-output pool as VMM memory - // (cuMem FABRIC/POSIX-FD via gpuCallocPhysical) so peers import TMA-eligible - // VAs. Falls back to the increment-3 inline carve (cudaIpc, no TMA) on - // platforms without fabric-IPC support. - static const bool ep_fabric_pool_ok = mscclpp::isNvlsSupported() && resolve_fabric_ipc_supported(); - if (ep_fabric_pool_ok) { - const size_t ep_pool_bytes = Config::recv_pool_bytes_static(num_ranks); - recv_pool_local_ptr_ = mscclpp::detail::gpuCallocPhysical(ep_pool_bytes); - CUDA_CHECK(cudaMemset(recv_pool_local_ptr_, 0, ep_pool_bytes)); - } else { - // inc3 fallback: append the pool AFTER the fifo/task regions of the NVL - // alloc (reached by peers via buffer_ptrs[peer] + recv_pool_off_), NOT in - // num_nvl_bytes so it does not count against the INT_MAX registered cap. - recv_pool_off_ = ((ep_nvl_alloc + 127) / 128) * 128; - ep_nvl_alloc = recv_pool_off_ + static_cast(Config::recv_pool_bytes_static(num_ranks)); - } -#endif - CUDA_CHECK(cudaMalloc(&buffer_ptrs[nvl_rank], ep_nvl_alloc)); - CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank])); - buffer_ptrs_gpu = - reinterpret_cast(reinterpret_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + fifo_bytes); - - // Set task fifo - EP_HOST_ASSERT(NUM_MAX_FIFO_SLOTS % num_nvl_ranks == 0); - task_fifo_ptrs[nvl_rank] = - reinterpret_cast(reinterpret_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes); - task_fifo_ptrs_gpu = reinterpret_cast(reinterpret_cast(buffer_ptrs[nvl_rank]) + num_nvl_bytes + - fifo_bytes + buffer_ptr_bytes); - - // No need to synchronize, will do a full device sync during `sync` - CUDA_CHECK(cudaMemsetAsync(task_fifo_ptrs[nvl_rank], 0, fifo_bytes, comm_stream)); - } - - // Create 32 MiB workspace - CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES)); - CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream)); - - // MoE counter - CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_counter_mapped, const_cast(moe_recv_counter), 0)); - *moe_recv_counter = -1; - - // MoE expert-level counter - CUDA_CHECK(cudaMallocHost(&moe_recv_expert_counter, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_expert_counter_mapped, const_cast(moe_recv_expert_counter), 0)); - for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++i) moe_recv_expert_counter[i] = -1; - - // MoE RDMA-level counter - if (num_rdma_ranks > 0) { - CUDA_CHECK(cudaMallocHost(&moe_recv_rdma_counter, sizeof(int), cudaHostAllocMapped)); - CUDA_CHECK(cudaHostGetDevicePointer(&moe_recv_rdma_counter_mapped, const_cast(moe_recv_rdma_counter), 0)); - *moe_recv_rdma_counter = -1; - } - - for (auto& ps : proxy_services) ps->startProxy(); + ringBufferBytes_ = config_.get_nvl_buffer_size_hint(static_cast(maxHiddenBytes_), numRanks_); + EP_HOST_ASSERT(ringBufferBytes_ <= static_cast(std::numeric_limits::max())); + taskFifoOffset_ = ringBufferBytes_; + symmetricBufferBytes_ = + configAlign(taskFifoOffset_ + sizeof(int) * NUM_MAX_FIFO_SLOTS, NUM_BUFFER_ALIGNMENT_BYTES); + physicalRingBuffer_ = numRanks_ > numNvlRanks_; + if (const char* direct = std::getenv("MSCCLPP_EP_INTRA_DIRECT"); direct != nullptr and std::atoi(direct) != 0) + recvPoolBytes_ = Config::recv_pool_bytes_static(numRanks_); + setup(communicator); } MoEHighThroughputRuntime::~MoEHighThroughputRuntime() noexcept(false) { - // Synchronize CUDA_CHECK(cudaDeviceSynchronize()); - - if (num_nvl_bytes > 0) { - // Barrier - intranode::barrier(task_fifo_ptrs_gpu, head, nvl_rank, num_nvl_ranks, comm_stream); - move_fifo_slots(); + if (available_) { + intranode::barrier(taskFifoPtrsGpu_, head_, rank_, numRanks_, nullptr); + moveFifoSlots(); CUDA_CHECK(cudaDeviceSynchronize()); - - // Close remote IPC - if (isAvailable()) { - for (int i = 0; i < num_nvl_ranks; ++i) - if (i != nvl_rank) CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i])); - } - - // Free local buffer and error flag - CUDA_CHECK(cudaFree(buffer_ptrs[nvl_rank])); - } -#ifdef EP_DISPATCH_NCCLEP - // Increment 4: release the device-side VMM recv-pool base array. The imported - // peer mappings are owned by recv_pool_remote_mems_ (auto-released); the local - // VMM pool follows the rdma_buffer_ptr convention (reclaimed at process exit). - if (recv_pool_ptrs_gpu != nullptr) { - CUDA_CHECK(cudaFree(recv_pool_ptrs_gpu)); - recv_pool_ptrs_gpu = nullptr; - } - if (recv_pool_global_ptrs_gpu != nullptr) { - CUDA_CHECK(cudaFree(recv_pool_global_ptrs_gpu)); - recv_pool_global_ptrs_gpu = nullptr; - } - if (ep_combine_recv_idx_gpu != nullptr) { - CUDA_CHECK(cudaFree(ep_combine_recv_idx_gpu)); - ep_combine_recv_idx_gpu = nullptr; - } -#endif - - // Free NVSHMEM - if (num_rdma_bytes > 0) { - // NVSHMEM support is not yet ported; if we got here with - // num_rdma_bytes > 0 the construction or sync call would already have - // failed, so there is nothing to tear down. } - // LL fast-path teardown. RegisteredMemory shared_ptrs in ll_memory_channels - // own the peer IPC mappings; we just release the device-side base array. - if (ll_ipc_ready) { - if (peer_rdma_bases_gpu != nullptr) { - CUDA_CHECK(cudaFree(peer_rdma_bases_gpu)); - peer_rdma_bases_gpu = nullptr; + if (combineRecvIdxGpu_ != nullptr) CUDA_CHECK(cudaFree(combineRecvIdxGpu_)); + if (recvPoolPtrsGpu_ != nullptr) CUDA_CHECK(cudaFree(recvPoolPtrsGpu_)); + if (taskFifoPtrsGpu_ != nullptr) CUDA_CHECK(cudaFree(taskFifoPtrsGpu_)); + if (bufferPtrsGpu_ != nullptr) CUDA_CHECK(cudaFree(bufferPtrsGpu_)); + if (moeRecvExpertCounter_ != nullptr) CUDA_CHECK(cudaFreeHost(const_cast(moeRecvExpertCounter_))); + if (moeRecvCounter_ != nullptr) CUDA_CHECK(cudaFreeHost(const_cast(moeRecvCounter_))); + + recvPoolMemories_.clear(); + peerMemories_.clear(); + if (recvPool_ != nullptr) mscclpp::detail::gpuFreePhysical(recvPool_); + if (symmetricBuffer_ != nullptr) { + if (physicalRingBuffer_) + mscclpp::detail::gpuFreePhysical(symmetricBuffer_); + else + CUDA_CHECK(cudaFree(symmetricBuffer_)); + } +} + +bool MoEHighThroughputRuntime::isAvailable() const { return available_; } + +bool MoEHighThroughputRuntime::isInternodeAvailable() const { return isAvailable() and numRanks_ > numNvlRanks_; } + +void MoEHighThroughputRuntime::setup(mscclpp::Communicator& communicator) { + EP_HOST_ASSERT(!available_); + if (physicalRingBuffer_) { + symmetricBuffer_ = mscclpp::detail::gpuCallocPhysical(symmetricBufferBytes_); + } else { + CUDA_CHECK(cudaMalloc(&symmetricBuffer_, symmetricBufferBytes_)); + CUDA_CHECK(cudaMemset(symmetricBuffer_, 0, symmetricBufferBytes_)); + } + if (recvPoolBytes_ > 0) recvPool_ = mscclpp::detail::gpuCallocPhysical(recvPoolBytes_); + + constexpr int RingBufferTag = 17; + constexpr int RecvPoolTag = 18; + const auto transport = mscclpp::Transport::CudaIpc; + peerMemories_.resize(numRanks_); + peerMemories_[rank_] = communicator.registerMemory(symmetricBuffer_, symmetricBufferBytes_, transport); + std::vector> remoteMemories(numRanks_); + std::vector> remoteRecvPools; + if (recvPool_ != nullptr) { + recvPoolMemories_.resize(numRanks_); + recvPoolMemories_[rank_] = communicator.registerMemory(recvPool_, recvPoolBytes_, transport); + remoteRecvPools.resize(numRanks_); + } + for (int peer = 0; peer < numRanks_; ++peer) { + if (peer == rank_) continue; + communicator.sendMemory(peerMemories_[rank_], peer, RingBufferTag); + remoteMemories[peer] = communicator.recvMemory(peer, RingBufferTag); + if (recvPool_ != nullptr) { + communicator.sendMemory(recvPoolMemories_[rank_], peer, RecvPoolTag); + remoteRecvPools[peer] = communicator.recvMemory(peer, RecvPoolTag); } } - for (auto& ps : proxy_services) ps->stopProxy(); - - // Free cuBLAS handle, workspace and MoE counter - CUDA_CHECK(cudaFree(workspace)); - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_counter))); - - // Free chunked mode staffs - CUDA_CHECK(cudaFreeHost(const_cast(moe_recv_expert_counter))); - - // Destroy the communication stream + ordering event. - if (comm_event) CUDA_CHECK(cudaEventDestroy(comm_event)); - if (comm_stream) CUDA_CHECK(cudaStreamDestroy(comm_stream)); -} - -void MoEHighThroughputRuntime::move_fifo_slots(int num_slots) { - head = (head + num_ranks * num_slots) % NUM_MAX_FIFO_SLOTS; -} - -void MoEHighThroughputRuntime::stream_wait(cudaStream_t dst, cudaStream_t src) { - // Make `dst` wait for everything currently enqueued on `src`. The reusable - // event is recorded on `src` then waited on `dst`; record-then-wait is - // host-sequential within a call so a single event is safe to reuse. - CUDA_CHECK(cudaEventRecord(comm_event, src)); - CUDA_CHECK(cudaStreamWaitEvent(dst, comm_event, 0)); -} - -bool MoEHighThroughputRuntime::isAvailable() const { return available; } - -bool MoEHighThroughputRuntime::isInternodeAvailable() const { return isAvailable() and num_ranks > NUM_MAX_NVL_PEERS; } - -int MoEHighThroughputRuntime::getNumRdmaRanks() const { return num_rdma_ranks; } - -int MoEHighThroughputRuntime::getRdmaRank() const { return rdma_rank; } - -int MoEHighThroughputRuntime::getRootRdmaRank(bool global) const { return global ? nvl_rank : 0; } - -int MoEHighThroughputRuntime::getLocalDeviceId() const { return device_id; } - -std::string MoEHighThroughputRuntime::getLocalIpcHandle() const { - return {ipc_handles[nvl_rank].reserved, CUDA_IPC_HANDLE_SIZE}; -} - -std::string MoEHighThroughputRuntime::getLocalNvshmemUniqueId() const { - // The MSCCL++ EP port replaces NVSHMEM with PortChannel/MemoryChannel, - // so there is no NVSHMEM unique id to expose. Kept for ABI parity with - // DeepEP's Python frontend; callers should use the MSCCL++ bootstrap. - throw std::runtime_error( - "mscclpp::ep::MoEHighThroughputRuntime::getLocalNvshmemUniqueId: not applicable (NVSHMEM is not used in " - "mscclpp_ep)"); -} - -mscclpp::UniqueId MoEHighThroughputRuntime::createUniqueId() const { return bootstrap->createUniqueId(); } - -void MoEHighThroughputRuntime::connect(mscclpp::UniqueId rootId) { - bootstrap->initialize(rootId); - communicator = std::make_shared(bootstrap); -} - -void MoEHighThroughputRuntime::sync(const std::vector& deviceIds, - const std::vector>& allGatheredHandles, - const std::optional& rootUniqueIdOpt) { - EP_HOST_ASSERT(not isAvailable()); - - const std::vector ib_transports = { - mscclpp::Transport::IB0, mscclpp::Transport::IB1, mscclpp::Transport::IB2, mscclpp::Transport::IB3, - mscclpp::Transport::IB4, mscclpp::Transport::IB5, mscclpp::Transport::IB6, mscclpp::Transport::IB7}; - const auto ipc_transport = mscclpp::Transport::CudaIpc; - const auto ib_transport = ib_transports[device_id]; - const mscclpp::TransportFlags all_transport = ipc_transport | ib_transport; - - // Sync IPC handles - if (num_nvl_bytes > 0) { - EP_HOST_ASSERT(num_ranks == deviceIds.size()); - EP_HOST_ASSERT(deviceIds.size() == allGatheredHandles.size()); - for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++i) { - EP_HOST_ASSERT(allGatheredHandles[offset + i].has_value()); - const auto& handle_str = allGatheredHandles[offset + i].value(); - EP_HOST_ASSERT(handle_str.size() == CUDA_IPC_HANDLE_SIZE); - if (offset + i != rank) { - std::memcpy(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE); - CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess)); - task_fifo_ptrs[i] = reinterpret_cast(reinterpret_cast(buffer_ptrs[i]) + num_nvl_bytes); - } else { - EP_HOST_ASSERT(std::memcmp(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE) == 0); - } + bufferPtrs_.resize(numRanks_); + taskFifoPtrs_.resize(numRanks_); + recvPoolPtrs_.resize(numRanks_); + for (int peer = 0; peer < numRanks_; ++peer) { + if (peer != rank_) { + peerMemories_[peer] = remoteMemories[peer].get(); + if (recvPool_ != nullptr) recvPoolMemories_[peer] = remoteRecvPools[peer].get(); } + void* base = peer == rank_ ? symmetricBuffer_ : peerMemories_[peer].data(); + bufferPtrs_[peer] = base; + taskFifoPtrs_[peer] = reinterpret_cast(static_cast(base) + taskFifoOffset_); + recvPoolPtrs_[peer] = recvPool_ == nullptr ? nullptr : (peer == rank_ ? recvPool_ : recvPoolMemories_[peer].data()); + } - // Copy all buffer and task pointers to GPU - CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, sizeof(void*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMalloc(&bufferPtrsGpu_, sizeof(void*) * numRanks_)); + CUDA_CHECK(cudaMalloc(&taskFifoPtrsGpu_, sizeof(int*) * numRanks_)); + CUDA_CHECK(cudaMemcpy(bufferPtrsGpu_, bufferPtrs_.data(), sizeof(void*) * numRanks_, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(taskFifoPtrsGpu_, taskFifoPtrs_.data(), sizeof(int*) * numRanks_, cudaMemcpyHostToDevice)); + if (recvPool_ != nullptr) { + CUDA_CHECK(cudaMalloc(&recvPoolPtrsGpu_, sizeof(void*) * numRanks_)); + CUDA_CHECK(cudaMemcpy(recvPoolPtrsGpu_, recvPoolPtrs_.data(), sizeof(void*) * numRanks_, cudaMemcpyHostToDevice)); CUDA_CHECK( - cudaMemcpy(task_fifo_ptrs_gpu, task_fifo_ptrs, sizeof(int*) * NUM_MAX_NVL_PEERS, cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaDeviceSynchronize()); - - // create connections - std::vector connections; - { - std::vector> connection_futures; - mscclpp::EndpointConfig local_config(ipc_transport); - for (int i = 0; i < num_nvl_ranks; ++i) { - auto r = i + rdma_rank * num_nvl_ranks; - connection_futures.emplace_back(communicator->connect(local_config, r, 0)); - } - for (auto& future : connection_futures) { - connections.emplace_back(future.get()); - } - } - - auto buffer_mem = communicator->registerMemory(buffer_ptrs[nvl_rank], num_nvl_bytes, ipc_transport); - - std::vector> remote_mem_futures(num_nvl_ranks); - for (int i = 0; i < num_nvl_ranks; ++i) { - if (i == nvl_rank) continue; - auto r = i + rdma_rank * num_nvl_ranks; - communicator->sendMemory(buffer_mem, r, 0); - remote_mem_futures[i] = communicator->recvMemory(r, 0); - } - for (int i = 0; i < num_nvl_ranks; ++i) { - if (i == nvl_rank) continue; - auto sema = std::make_shared(*communicator, connections[i]); - memory_channels.emplace_back(sema, remote_mem_futures[i].get(), buffer_mem); - } - std::vector memory_channel_handles(num_nvl_ranks); - for (int i = 0; i < num_nvl_ranks; ++i) { - if (i == nvl_rank) continue; - memory_channel_handles[i] = memory_channels.rbegin()->deviceHandle(); - } - -#ifdef EP_DISPATCH_NCCLEP - // Increment 4: exchange the VMM recv-output pool base pointers across the - // local NVL peers. registerMemory(ipc) on a gpuCallocPhysical buffer carries - // FABRIC/POSIX-FD handles, so the imported `.data()` peer VA is dereferenceable - // (and TMA-eligible) — unlike the cudaIpc-mapped buffer_ptrs. Imported - // RegisteredMemory is retained in recv_pool_remote_mems_ to keep the mapping. - if (recv_pool_local_ptr_ != nullptr) { - constexpr int kRecvPoolTag = 7; - const size_t ep_pool_bytes = Config::recv_pool_bytes_static(num_ranks); - auto pool_mem = communicator->registerMemory(recv_pool_local_ptr_, ep_pool_bytes, ipc_transport); - std::vector> pool_futs(num_nvl_ranks); - for (int i = 0; i < num_nvl_ranks; ++i) { - if (i == nvl_rank) continue; - auto r = i + rdma_rank * num_nvl_ranks; - communicator->sendMemory(pool_mem, r, kRecvPoolTag); - pool_futs[i] = communicator->recvMemory(r, kRecvPoolTag); - } - recv_pool_ptrs_.assign(NUM_MAX_NVL_PEERS, nullptr); - recv_pool_ptrs_[nvl_rank] = recv_pool_local_ptr_; - recv_pool_remote_mems_.resize(num_nvl_ranks); - for (int i = 0; i < num_nvl_ranks; ++i) { - if (i == nvl_rank) continue; - recv_pool_remote_mems_[i] = pool_futs[i].get(); - recv_pool_ptrs_[i] = recv_pool_remote_mems_[i].data(); - } - CUDA_CHECK(cudaMalloc(&recv_pool_ptrs_gpu, sizeof(void*) * NUM_MAX_NVL_PEERS)); - CUDA_CHECK(cudaMemcpy(recv_pool_ptrs_gpu, recv_pool_ptrs_.data(), sizeof(void*) * NUM_MAX_NVL_PEERS, - cudaMemcpyHostToDevice)); - // Intranode combine-direct: allocate the per-(token, dst rank) gather map so the - // INTRA_DIRECT dispatch sender can record each token's recv-pool slot and the TMA - // direct-gather combine can read it back. Idempotent with the internode (inc5) - // allocation below (guarded by == nullptr); sized for the worst-case source tokens. - if (ep_combine_recv_idx_gpu == nullptr) - CUDA_CHECK(cudaMalloc(&ep_combine_recv_idx_gpu, - sizeof(int) * static_cast(Config::kEpRecvPoolMaxTokens) * num_ranks)); - - if (rank == 0) { - printf("[mscclpp_ep] inc4 VMM recv-pool peer bases (rank 0):"); - for (int i = 0; i < NUM_MAX_NVL_PEERS; ++i) printf(" [%d]=%p", i, recv_pool_ptrs_[i]); - printf("\n"); - fflush(stdout); - } - } -#endif - - memory_channel_handles_device_ptr = - mscclpp::detail::gpuCallocShared(num_nvl_ranks); - mscclpp::gpuMemcpy( - memory_channel_handles_device_ptr.get(), memory_channel_handles.data(), num_nvl_ranks, cudaMemcpyHostToDevice); + cudaMalloc(&combineRecvIdxGpu_, sizeof(int) * static_cast(Config::RecvPoolMaxTokens) * numRanks_)); } - - // RDMA buffer setup (replaces DeepEP's NVSHMEM symmetric-heap allocation). - // - // Unlike DeepEP which used `nvshmem_align` to place the RDMA buffer on the - // symmetric heap, all our internode communication goes through MSCCL++ - // `PortChannel` (proxy-based RDMA), so a plain `cudaMalloc` + IB memory - // registration is sufficient. The bootstrap barrier replaces - // `nvshmem_barrier_all`. - if (num_rdma_bytes > 0) { - EP_HOST_ASSERT(communicator != nullptr); - EP_HOST_ASSERT(bootstrap != nullptr); - - // Allocate the RDMA buffer. - // - // For low-latency mode on platforms that support NVLink-SHARP / NVLS - // (GB200 NVL72 with nvidia-imex configured), use mscclpp's - // `gpuCallocPhysical` (cuMemCreate + cuMemMap with POSIX_FD|FABRIC - // handle types) so the buffer is eligible for cuMem fabric IPC — the - // LL fast path then maps the buffer across the NVL72 fabric via - // nvidia-imex and performs atomicAdd over NVLink rather than RDMA - // (which has IBV_ATOMIC_NONE on Azure CX-7 RoCE). - // - // Fallback: on platforms without NVLS / multicast support (e.g. - // H100 + IB, A100 + IB), `gpuCallocPhysical` would either fail or - // produce non-fabric-IPC memory; fall back to plain `cudaMalloc` and - // let the LL path use the existing PortChannel proxy mechanism over - // IB. For HT internode on Azure GB200 (Phase 4), we also use - // fabric-IPC so cross-node `handle.put(data)` can be replaced by - // direct kernel-side writes via NVL72 fabric pointers — bypassing - // the broken Azure CX-7 RoCE RDMA WRITE path. - // - // NVLS-supported but FABRIC-unsupported deployments (e.g. H100 on - // Azure NDv5) must not take this path for the HT cross-host case: - // their `gpuCallocPhysical` result is a POSIX-FD-only handle whose - // cross-host import falls back to a master-local unix socket which - // worker ranks cannot reach. - static const bool fabric_ipc_supported = resolve_fabric_ipc_supported(); - const bool use_fabric_ipc_alloc = - mscclpp::isNvlsSupported() && fabric_ipc_supported && (low_latency_mode || num_rdma_ranks > 1); - if (use_fabric_ipc_alloc) { - rdma_buffer_ptr = mscclpp::detail::gpuCallocPhysical(num_rdma_bytes); - CUDA_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes)); - } else { - CUDA_CHECK(cudaMalloc(&rdma_buffer_ptr, num_rdma_bytes)); - CUDA_CHECK(cudaMemset(rdma_buffer_ptr, 0, num_rdma_bytes)); - } - if (rank == 0) { - printf("[mscclpp_ep] rdma_buffer allocator: %s (low_latency=%d, nvls=%d, fabric_ipc=%d)\n", - use_fabric_ipc_alloc ? "gpuCallocPhysical (fabric-IPC)" : "cudaMalloc", (int)low_latency_mode, - (int)mscclpp::isNvlsSupported(), (int)fabric_ipc_supported); - fflush(stdout); - } - bootstrap->barrier(); - CUDA_CHECK(cudaDeviceSynchronize()); - - // Rank -> RDMA buffer IDs. MemoryIds are local to each ProxyService; - // we register every memory in every proxy in the same global order so - // a single int identifies the memory across all of them. - std::map memory_ids; - - auto add_memory_to_all = [&](mscclpp::RegisteredMemory mem) -> mscclpp::MemoryId { - mscclpp::MemoryId id = static_cast(-1); - for (auto& ps : proxy_services) { - auto cur = ps->addMemory(mem); - if (id == static_cast(-1)) id = cur; - EP_HOST_ASSERT(cur == id && "MemoryIds drifted across proxy services"); - } - return id; - }; - - // Register local memory - auto local_rdma_buffer_mem = communicator->registerMemory(rdma_buffer_ptr, num_rdma_bytes, all_transport); - memory_ids[rank] = add_memory_to_all(local_rdma_buffer_mem); - - // Send local memory to other ranks. - // - // NOTE: DeepEP filters this to same-GPU-ID peers in low_latency_mode - // because LL there uses NVSHMEM, not port channels. This port drives - // LL kernels through PortChannel, so every peer must have a real - // memory/connection/semaphore/port channel entry. Treat LL and HT - // sync identically: always connect all peers. - // - // Caveat: for a pure intra-node LL launch (``num_nvl_bytes == 0`` with - // every peer on the same host) the resulting port channels go through - // the CPU proxy over IB loopback between different HCAs, which on - // this platform does not deliver atomics reliably and currently - // deadlocks LL dispatch. See `src/ext/ep/README.md` for the full - // discussion. Cross-node LL (DeepEP's recommended 1-GPU-per-node - // topology) is unaffected. - // Use tag=1 to disambiguate from the NVL phase's tag=0 traffic with same-node peers. - constexpr int kRdmaTag = 1; - for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; - communicator->sendMemory(local_rdma_buffer_mem, r, kRdmaTag); - } - - // Receive remote memory from other ranks. - for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; - auto f = communicator->recvMemory(r, kRdmaTag); - auto mem = f.get(); - memory_ids[r] = add_memory_to_all(std::move(mem)); - } - - // Rank -> vector of connections - std::unordered_map> connections; - const mscclpp::EndpointConfig ipc_cfg(ipc_transport); - const mscclpp::EndpointConfig ib_cfg(ib_transport); - - // Self connection for local memory (CUDA IPC). - connections[rank].emplace_back(communicator->connect(ipc_cfg, rank, kRdmaTag).get()); - - // Remote IB connections (multi-QP per peer). - const int num_ib_connections_per_rank = 12; // #QPs per rank (mirrors DeepEP). - for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; - std::vector> futures; - futures.reserve(num_ib_connections_per_rank); - for (int i = 0; i < num_ib_connections_per_rank; ++i) { - futures.emplace_back(communicator->connect(ib_cfg, r, kRdmaTag)); - } - for (auto& f : futures) connections[r].emplace_back(f.get()); - } - - // Rank -> vector of (proxy_idx, semaphore_id_within_proxy). Iterate - // peers in sorted rank order so semaphore pairings between nodes line - // up deterministically. Channels — and therefore their backing - // semaphores — are sharded across `proxy_services`: channel at flat - // index `i*num_ranks + r` lives on proxy `(i*num_ranks + r) % - // num_proxy_services`. SemaphoreIds are local to each proxy, so we - // record (proxy_idx, sid) pairs. - std::unordered_map>> sema_ids; - const int num_semaphores_per_rank = 16; - for (int i = 0; i < num_semaphores_per_rank; ++i) { - for (int r = 0; r < num_ranks; ++r) { - auto conn_it = connections.find(r); - EP_HOST_ASSERT(conn_it != connections.end()); - auto& conns = conn_it->second; - auto& conn = conns[i % conns.size()]; - int proxy_idx = (i * num_ranks + r) % num_proxy_services; - auto sema_id = proxy_services[proxy_idx]->buildAndAddSemaphore(*communicator, conn); - sema_ids[r].emplace_back(proxy_idx, sema_id); - } - } - - // Create port channels + device handles. - // - // The kernels index `port_channel_handles[channel_id * num_ranks + peer_rank]` - // where peer_rank is a GLOBAL rank in [0..num_ranks). So the outer stride must - // be num_ranks with peers in ascending rank order. Iterating `memory_ids` (an - // `unordered_map`) yields hash order and would misroute signals, deadlocking. - // Each channel inherits the proxy of the semaphore it was built on, so the - // resulting `PortChannelDeviceHandle` routes its FIFO pushes to the correct - // proxy thread. - const int num_port_channels_per_rank = num_semaphores_per_rank; - std::vector port_channel_handles; - for (int i = 0; i < num_port_channels_per_rank; ++i) { - for (int r = 0; r < num_ranks; ++r) { - auto mem_it = memory_ids.find(r); - EP_HOST_ASSERT(mem_it != memory_ids.end()); - auto memory_id = mem_it->second; - auto [proxy_idx, sema_id] = sema_ids[r][i % sema_ids[r].size()]; - auto port_channel = proxy_services[proxy_idx]->portChannel(sema_id, memory_id, memory_ids[rank]); - port_channels.emplace_back(std::move(port_channel)); - port_channel_handles.emplace_back(port_channels.rbegin()->deviceHandle()); - } - } - - port_channel_handles_device_ptr = - mscclpp::detail::gpuCallocShared(port_channel_handles.size()); - mscclpp::gpuMemcpy(port_channel_handles_device_ptr.get(), - port_channel_handles.data(), port_channel_handles.size(), - cudaMemcpyHostToDevice); - - // ------------------------------------------------------------------ - // HT internode NVLS multicast setup (Wide Proposal B2). - // - // On platforms with NVLink-SHARP / multicast support (GB200 NVL72 - // with nvidia-imex), allocate a multicast-bound buffer used by the HT - // dispatch / combine / notify_dispatch kernels for: - // - tail / head atomic counters (replaces the 4 atomicAdd sites) - // - notify_dispatch barrier epoch (replaces port_channel.signal/wait) - // - notify_dispatch small-data delivery (replaces putWithSignal) - // - // All cross-node atomic adds become `multimem.red.add.u64` PTX which - // travels over the NVL72 fabric — bypassing IB entirely. This is the - // same fabric path that LL Proposal A validated cross-node. - // - // Fallback (existing IB platforms): when `isNvlsSupported()` is false - // or there is only one RDMA rank (intranode-only), `nvls_ht_enabled` - // stays `false` and kernels select the legacy PortChannel path. - // - // Additionally: NVLS alone is insufficient for cross-host. H100 has - // NVLS within a node but no cross-host NVSwitch fabric, so - // `connectNvlsCollective` either fails or builds a per-node-only - // multicast object, and the POSIX-FD fallback for handle exchange - // routes through a master-local unix socket that worker ranks cannot - // reach. Gate this path on `fabric_ipc_supported` so non-Blackwell - // deployments cleanly use the legacy PortChannel path. - // - // Skipped for `low_latency_mode` since LL has its own (working) - // fabric-IPC path via Proposal A and does not use the HT counter - // protocol. - // ------------------------------------------------------------------ - nvls_ht_enabled = false; - if (!low_latency_mode && num_rdma_ranks > 1 && mscclpp::isNvlsSupported() && fabric_ipc_supported) { - // Worst-case sizing — chosen so the same multicast buffer fits any - // (num_sms, num_rdma_ranks) configuration the kernels may launch with. - const size_t kCounterBytesPerChannel = - static_cast(NUM_MAX_RDMA_PEERS) * NUM_MAX_RDMA_PEERS * sizeof(uint64_t); - const size_t tail_bytes = static_cast(kNvlsMaxChannels) * kCounterBytesPerChannel; - const size_t head_bytes = tail_bytes; - const size_t barrier_bytes = static_cast(kNvlsBarrierSlots) * sizeof(uint64_t); - // Data region: per-sender slot of [num_rdma_ranks_max × per_peer_bytes], - // one slot per global rank (worst-case num_ranks = NUM_MAX_RDMA_PEERS * - // NUM_MAX_NVL_PEERS). Each rank writes its own slot via `multimem.st`; - // every receiver then reads the sub-position destined to it. - const size_t kPerSenderSlotBytes = static_cast(NUM_MAX_RDMA_PEERS) * kNvlsPerPeerBytes; - const size_t kMaxRanks = static_cast(NUM_MAX_RDMA_PEERS) * NUM_MAX_NVL_PEERS; - const size_t data_bytes = kMaxRanks * kPerSenderSlotBytes; - - // 256 B alignment for each sub-region to keep `multimem` ops well-aligned. - auto align256 = [](size_t x) { return (x + 255) & ~size_t(255); }; - nvls_ht_off_tail = 0; - nvls_ht_off_head = align256(nvls_ht_off_tail + tail_bytes); - nvls_ht_off_barrier = align256(nvls_ht_off_head + head_bytes); - nvls_ht_off_data = align256(nvls_ht_off_barrier + barrier_bytes); - nvls_ht_total_bytes = align256(nvls_ht_off_data + data_bytes); - - // GpuBuffer auto-uses gpuCallocPhysicalShared (cuMem fabric handle) - // when isNvlsSupported() — required for multicast bind. - nvls_ht_buffer = std::make_shared>(nvls_ht_total_bytes); - CUDA_CHECK(cudaMemset(nvls_ht_buffer->data(), 0, nvls_ht_buffer->bytes())); - - std::vector all_ranks; - all_ranks.reserve(num_ranks); - for (int r = 0; r < num_ranks; ++r) all_ranks.push_back(r); - - // Collective: every rank must call. If it fails (e.g. IMEX - // misconfigured, peers in different fabrics), the exception - // propagates — there is no clean fallback mid-collective. The - // `isNvlsSupported()` gate above is the production guard. - nvls_ht_conn = mscclpp::connectNvlsCollective(communicator, all_ranks, nvls_ht_buffer->bytes()); - auto sw = nvls_ht_conn->bindAllocatedMemory(reinterpret_cast(nvls_ht_buffer->data()), - nvls_ht_buffer->bytes()); - nvls_ht_sc = std::make_shared(std::move(sw)); - auto h = nvls_ht_sc->deviceHandle(); - nvls_ht_mc_ptr = h.mcPtr; - nvls_ht_dev_ptr = h.devicePtr; - nvls_ht_enabled = (nvls_ht_mc_ptr != nullptr) && (nvls_ht_dev_ptr != nullptr); - - // DIAG: print mcPtr/devicePtr/buf-VA per rank to verify whether - // connectNvlsCollective produced a multicast that actually spans - // both nodes (suspected: per-node only on Azure GB200). - printf("[mscclpp_ep] NVLS HT diag rank=%d mcPtr=%p devicePtr=%p bufVA=%p bytes=%zu\n", rank, - (void*)nvls_ht_mc_ptr, (void*)nvls_ht_dev_ptr, (void*)nvls_ht_buffer->data(), - (size_t)nvls_ht_buffer->bytes()); - fflush(stdout); - - bootstrap->barrier(); - - if (rank == 0) { - printf( - "[mscclpp_ep] NVLS HT multicast: enabled=%d total=%zu KB " - "(tail@%zu head@%zu barrier@%zu data@%zu)\n", - (int)nvls_ht_enabled, nvls_ht_total_bytes / 1024, nvls_ht_off_tail, nvls_ht_off_head, nvls_ht_off_barrier, - nvls_ht_off_data); - fflush(stdout); - } - } else if (rank == 0) { - printf( - "[mscclpp_ep] NVLS HT multicast: disabled (low_latency=%d, num_rdma_ranks=%d, " - "nvls_supported=%d, fabric_ipc=%d)\n", - (int)low_latency_mode, num_rdma_ranks, (int)mscclpp::isNvlsSupported(), (int)fabric_ipc_supported); - fflush(stdout); - } - - // ------------------------------------------------------------------ - // Intra-node LL fast path setup. - // - // When all ranks sit on the same host (num_rdma_ranks == 1), LL dispatch - // and combine still go through `PortChannel` above — which internally - // uses the proxy service over IB loopback between different HCAs on - // this platform. That path is correct but slow (caps at ~170 GB/s vs. - // NVLink's multi-TB/s). We additionally set up CUDA-IPC peer pointers - // to each peer's `rdma_buffer_ptr` plus a set of per-peer MemoryChannels - // for a barrier ring. The LL kernels select this path at launch time. - // Cross-node LL uses cuMem fabric IPC (Proposal A): peers map - // `rdma_buffer_ptr` through the NVL72 NVSwitch fabric via nvidia-imex, - // and the LL kernels do direct `st.global` + atomicAdd through those - // peer pointers. This bypasses the RDMA path entirely (Azure CX-7 RoCE - // has IBV_ATOMIC_NONE which makes the proxy-emulated atomicAdd hang). - // Requires nvidia-imex active on every rank's host with a shared - // `nodes_config.cfg` covering all node IPs. - // ------------------------------------------------------------------ - // ------------------------------------------------------------------ - // Single-node LL uses regular CUDA IPC peer pointers. Cross-node LL/HT on - // GB200 uses fabric-IPC peer pointers, letting kernels write tokens directly - // into each peer's `rdma_buffer_ptr` via the NVL72 fabric VA. - // ------------------------------------------------------------------ - const int ipc_domain_size = use_fabric_ipc_alloc ? num_ranks : num_nvl_ranks; - auto is_ipc_peer = [&](int peer) { - return peer != rank && ipc_domain_size > 1 && rank / ipc_domain_size == peer / ipc_domain_size; - }; - const bool want_peer_ipc = use_fabric_ipc_alloc || (low_latency_mode && ipc_domain_size > 1); - if (want_peer_ipc) { - // Reuse the local RDMA registration's CudaIpc transport entry. The - // existing `local_rdma_buffer_mem` was registered with `all_transport` - // (= ipc | ib), so its CudaIpc TransportInfo is already populated - // with the FABRIC handle (when supported by the underlying physical - // allocation). We need a separate registration only because the - // remote-side `recvMemory` below is tagged independently. - constexpr int kLlIpcTag = 2; - auto rdma_mem_ipc = communicator->registerMemory(rdma_buffer_ptr, num_rdma_bytes, ipc_transport); - std::vector> remote_futures(num_ranks); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank || !is_ipc_peer(r)) continue; - communicator->sendMemory(rdma_mem_ipc, r, kLlIpcTag); - remote_futures[r] = communicator->recvMemory(r, kLlIpcTag); - } - std::vector ll_ipc_conns(num_ranks); - { - std::vector> conn_futures(num_ranks); - mscclpp::EndpointConfig cfg(ipc_transport); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank || !is_ipc_peer(r)) continue; - conn_futures[r] = communicator->connect(cfg, r, kLlIpcTag); - } - for (int r = 0; r < num_ranks; ++r) { - if (r == rank || !is_ipc_peer(r)) continue; - ll_ipc_conns[r] = conn_futures[r].get(); - } - } - - // Resolve peer base pointers from the (now mapped) remote - // RegisteredMemory. `data()` returns the locally-mapped peer pointer; - // for fabric handles this address lives in the cuMem fabric VA range, - // while single-node LL gets a regular CUDA IPC mapping. - peer_rdma_bases.assign(num_ranks, nullptr); - peer_rdma_bases[rank] = rdma_buffer_ptr; - std::vector remote_mems(num_ranks); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank || !is_ipc_peer(r)) continue; - remote_mems[r] = remote_futures[r].get(); - peer_rdma_bases[r] = remote_mems[r].data(); - } - CUDA_CHECK(cudaMalloc(&peer_rdma_bases_gpu, sizeof(void*) * num_ranks)); - CUDA_CHECK( - cudaMemcpy(peer_rdma_bases_gpu, peer_rdma_bases.data(), sizeof(void*) * num_ranks, cudaMemcpyHostToDevice)); - if (rank == 0) { - printf("[mscclpp_ep] Phase 4 fabric-IPC peer bases (rank 0):\n"); - for (int r = 0; r < num_ranks; ++r) { - printf(" peer_rdma_bases[%d] = %p\n", r, peer_rdma_bases[r]); - } - fflush(stdout); - } - -#ifdef EP_DISPATCH_NCCLEP - { - // Increment 5 (inc5 flat-domain dispatch): when MSCCLPP_EP_DIRECT is set, - // exchange the cuMem-fabric recv-output pool base to ALL ranks (indexed by - // global rank), mirroring the peer_rdma_bases exchange. Lets the RDMA sender - // write each token directly into the destination GPU's recv pool over - // fabric-VA, removing the rail-aligned rdma_channel bounce + forwarder - // transpose. Same cuMem FABRIC handle as inc4a's pool. Env-gated so the - // inc4a baseline (MSCCLPP_EP_DIRECT unset) is byte-for-byte unchanged. - const char* e_direct = std::getenv("MSCCLPP_EP_DIRECT"); - const bool ep_direct = (e_direct != nullptr && std::atoi(e_direct) != 0); - if (ep_direct && recv_pool_local_ptr_ != nullptr) { - constexpr int kRecvPoolGlobalTag = 9; - const size_t ep_pool_bytes_g = Config::recv_pool_bytes_static(num_ranks); - auto pool_mem_g = communicator->registerMemory(recv_pool_local_ptr_, ep_pool_bytes_g, all_transport); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; - communicator->sendMemory(pool_mem_g, r, kRecvPoolGlobalTag); - } - std::vector> gpool_futs(num_ranks); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; - gpool_futs[r] = communicator->recvMemory(r, kRecvPoolGlobalTag); - } - recv_pool_global_ptrs_.assign(num_ranks, nullptr); - recv_pool_global_ptrs_[rank] = recv_pool_local_ptr_; - recv_pool_global_remote_mems_.resize(num_ranks); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; - recv_pool_global_remote_mems_[r] = gpool_futs[r].get(); - recv_pool_global_ptrs_[r] = recv_pool_global_remote_mems_[r].data(); - } - CUDA_CHECK(cudaMalloc(&recv_pool_global_ptrs_gpu, sizeof(void*) * num_ranks)); - CUDA_CHECK(cudaMemcpy(recv_pool_global_ptrs_gpu, recv_pool_global_ptrs_.data(), sizeof(void*) * num_ranks, - cudaMemcpyHostToDevice)); - if (rank == 0) { - printf("[mscclpp_ep] inc5 domain-wide recv-pool bases (rank 0):"); - for (int r = 0; r < num_ranks; ++r) printf(" [%d]=%p", r, recv_pool_global_ptrs_[r]); - printf("\n"); - fflush(stdout); - } - // inc5 combine-direct (Stage 1): allocate the per-(token, dst global - // rank) gather map once, sized for the worst-case source token count. - // The dispatch sender fills it; combine gathers from it. - if (ep_combine_recv_idx_gpu == nullptr) - CUDA_CHECK(cudaMalloc(&ep_combine_recv_idx_gpu, - sizeof(int) * static_cast(Config::kEpRecvPoolMaxTokens) * num_ranks)); - } - } -#endif - - if (low_latency_mode) { - // LL barrier ring needs MemoryChannels. - std::vector ll_handles(num_ranks); - for (int r = 0; r < num_ranks; ++r) { - if (r == rank || !is_ipc_peer(r)) continue; - auto sema = std::make_shared(*communicator, ll_ipc_conns[r]); - ll_memory_channels.emplace_back(sema, remote_mems[r], rdma_mem_ipc); - ll_handles[r] = ll_memory_channels.rbegin()->deviceHandle(); - } - ll_memory_channel_handles_device_ptr = - mscclpp::detail::gpuCallocShared(num_ranks); - mscclpp::gpuMemcpy( - ll_memory_channel_handles_device_ptr.get(), ll_handles.data(), num_ranks, cudaMemcpyHostToDevice); - - ll_ranks_per_ipc_domain = ipc_domain_size; - ll_ipc_ready = ipc_domain_size >= num_ranks; - } - } - } - - // Ready to use - available = true; + CUDA_CHECK(cudaMallocHost(&moeRecvCounter_, sizeof(int), cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moeRecvCounterMapped_, const_cast(moeRecvCounter_), 0)); + CUDA_CHECK(cudaMallocHost(&moeRecvExpertCounter_, sizeof(int) * NUM_MAX_LOCAL_EXPERTS, cudaHostAllocMapped)); + CUDA_CHECK(cudaHostGetDevicePointer(&moeRecvExpertCounterMapped_, const_cast(moeRecvExpertCounter_), 0)); + *moeRecvCounter_ = -1; + for (int i = 0; i < NUM_MAX_LOCAL_EXPERTS; ++i) moeRecvExpertCounter_[i] = -1; + available_ = true; } -void MoEHighThroughputRuntime::getDispatchLayout(int* numTokensPerRank, int* numTokensPerRdmaRank, - int* numTokensPerExpert, bool* isTokenInRank, const int64_t* topkIdx, - int numTokens, int numTopk, int numExperts, cudaStream_t stream) { - EP_HOST_ASSERT(numExperts > 0); - - // Make comm_stream wait for the caller's stream (replaces the torch - // stream_wait dance), then launch the layout kernel on comm_stream. - stream_wait(comm_stream, stream); - - internode::get_dispatch_layout(topkIdx, numTokensPerRank, numTokensPerRdmaRank, numTokensPerExpert, isTokenInRank, - numTokens, numTopk, num_ranks, numExperts, comm_stream); - - // Make the caller's stream wait for comm_stream so the outputs are visible. - stream_wait(stream, comm_stream); +void MoEHighThroughputRuntime::moveFifoSlots(int numSlots) { + head_ = (head_ + numRanks_ * numSlots) % NUM_MAX_FIFO_SLOTS; } -void MoEHighThroughputRuntime::computeIntranodeChannels(int xElementSize, const Config& config, int& dispatchNumSms, - bool& allSender, int& numChannels) const { - // MSCCLPP_EP_DISPATCH_NSM (intranode): set the dispatch block count independently of the - // combine grid. The dispatch is channel-partitioned (notify + senders share num_channels), - // so DISPATCH_NSM drives this whole dispatch+layout pipeline consistently. Constrained to - // [2, config.num_sms] (the NVL buffers are sized for config.num_sms/2 channels), even. - int dispatch_num_sms = config.num_sms; - { - const char* e_dnsm = std::getenv("MSCCLPP_EP_DISPATCH_NSM"); - if (e_dnsm != nullptr) { - int n = std::atoi(e_dnsm); - n &= ~1; // round down to even (two blocks per channel) - if (n >= 2) dispatch_num_sms = std::min(n, config.num_sms); - } +void MoEHighThroughputRuntime::computeDispatchChannels(int xElementSize, int& dispatchNumSms, bool& allSender, + int& numChannels) const { + int dispatch_num_sms = config_.num_sms; + if (const char* env = std::getenv("MSCCLPP_EP_DISPATCH_NSM")) { + int value = std::atoi(env) & ~1; + if (value >= 2) dispatch_num_sms = std::min(value, config_.num_sms); } - // MSCCLPP_EP_INTRA_ALLSENDER (INTRA_DIRECT only): make EVERY block a sender (one channel - // per block) instead of the 50/50 sender/receiver split, so all dispatch_num_sms blocks - // move hidden directly to the dest pools (matching NCCL-EP's all-sender block count). - // Metadata goes to the pool META region + a drain pass. The all-sender layout is only - // consumable by the TMA combine (the ring fallback expects num_sms/2 channels), so this - // defaults ON under INTRA_DIRECT + TMA combine and auto-disables if the ring combine is - // forced (MSCCLPP_EP_COMBINE_TMA=0). Explicit MSCCLPP_EP_INTRA_ALLSENDER=0/1 overrides. - static const bool ep_intra_allsender_env = [] { - const char* d = std::getenv("MSCCLPP_EP_INTRA_DIRECT"); - if (not(d != nullptr and std::atoi(d) != 0)) return false; // requires INTRA_DIRECT - const char* ct = std::getenv("MSCCLPP_EP_COMBINE_TMA"); - if (ct != nullptr and std::atoi(ct) == 0) return false; // ring combine can't consume all-sender layout - const char* e = std::getenv("MSCCLPP_EP_INTRA_ALLSENDER"); - return e == nullptr or std::atoi(e) != 0; // default ON, disable only with explicit =0 + + static const bool AllSenderEnabled = [] { + const char* direct = std::getenv("MSCCLPP_EP_INTRA_DIRECT"); + if (not(direct != nullptr and std::atoi(direct) != 0)) return false; + const char* tma = std::getenv("MSCCLPP_EP_COMBINE_TMA"); + if (tma != nullptr and std::atoi(tma) == 0) return false; + const char* all_sender = std::getenv("MSCCLPP_EP_INTRA_ALLSENDER"); + return all_sender == nullptr or std::atoi(all_sender) != 0; }(); - // The all-sender path requires the BF16 direct pool (xElementSize == 2 mirrors the original - // `x.scalar_type() == torch::kBFloat16` gate). - const bool all_sender = ep_intra_allsender_env and recv_pool_local_ptr_ != nullptr and - recv_pool_ptrs_gpu != nullptr and xElementSize == 2; - // All-sender: one channel per block (num_channels = grid). Otherwise: two blocks per channel. + dispatchNumSms = dispatch_num_sms; - allSender = all_sender; - numChannels = all_sender ? dispatch_num_sms : dispatch_num_sms / 2; + allSender = AllSenderEnabled and xElementSize == 2; + numChannels = allSender ? dispatch_num_sms : dispatch_num_sms / 2; } -int MoEHighThroughputRuntime::getIntranodeDispatchNumChannels(int xElementSize, const Config& config) const { - int dispatch_num_sms = 0, num_channels = 0; +bool MoEHighThroughputRuntime::canUseDirectRecvPool(int numTokens, int numRecvTokens, int hidden, + int xElementSize) const { + const char* env = std::getenv("MSCCLPP_EP_INTRA_DIRECT"); + if (recvPool_ == nullptr or recvPoolPtrsGpu_ == nullptr or combineRecvIdxGpu_ == nullptr or env == nullptr or + std::atoi(env) == 0 or numTokens < 0 or numRecvTokens < 0 or hidden <= 0 or xElementSize <= 0 or + numTokens > Config::RecvPoolMaxTokens or numRecvTokens > Config::RecvPoolMaxTokens) + return false; + const int64_t hidden_bytes = static_cast(hidden) * xElementSize; + return hidden_bytes <= maxHiddenBytes_ and static_cast(numRecvTokens) * static_cast(hidden_bytes) <= + Config::get_recv_pool_hidden_bytes(numRanks_); +} + +void MoEHighThroughputRuntime::layout(int* numTokensPerRank, int* numTokensPerExpert, bool* isTokenInRank, + const int64_t* topkIdx, int numTokens, int numTopk, int numExperts, + cudaStream_t stream) { + EP_HOST_ASSERT(available_); + EP_HOST_ASSERT(numExperts > 0 and numExperts % numRanks_ == 0); + EP_HOST_ASSERT(numTopk > 0 and numTopk <= 32); + intranode::get_dispatch_layout(topkIdx, numTokensPerRank, numTokensPerExpert, isTokenInRank, numTokens, numTopk, + numRanks_, numExperts, stream); +} + +int MoEHighThroughputRuntime::getDispatchNumChannels(int xElementSize) const { + int dispatch_num_sms = 0; + int num_channels = 0; bool all_sender = false; - computeIntranodeChannels(xElementSize, config, dispatch_num_sms, all_sender, num_channels); + computeDispatchChannels(xElementSize, dispatch_num_sms, all_sender, num_channels); return num_channels; } -void* MoEHighThroughputRuntime::resolveIntranodeRecvXBuffer(int numRecvTokens, int hidden, int xElementSize, - const Config& config) const { - // Sender-direct (MSCCLPP_EP_INTRA_DIRECT): make recv_x a zero-copy view of this rank's - // peer-mapped recv pool so the sender writes hidden straight to its final slot, - // eliminating the 2-hop ring + receiver hidden drain. Mirrors the ep_intra_direct gate - // in intranodeDispatch exactly; returns nullptr when the caller should allocate recvX. -#ifdef EP_DISPATCH_NCCLEP - const size_t ep_intra_pool_header_bytes = config.get_recv_pool_header_bytes(num_ranks); - const char* e_intra_direct = std::getenv("MSCCLPP_EP_INTRA_DIRECT"); - const bool ep_intra_direct = e_intra_direct != nullptr && std::atoi(e_intra_direct) != 0 && - recv_pool_local_ptr_ != nullptr && recv_pool_ptrs_gpu != nullptr && - numRecvTokens <= Config::kEpRecvPoolMaxTokens && - static_cast(numRecvTokens) * hidden * static_cast(xElementSize) <= - static_cast(Config::recv_pool_bytes_static(num_ranks)) - - static_cast(ep_intra_pool_header_bytes); - if (ep_intra_direct) { - return static_cast(recv_pool_local_ptr_) + ep_intra_pool_header_bytes; - } - return nullptr; -#else - (void)numRecvTokens; - (void)hidden; - (void)xElementSize; - (void)config; - return nullptr; -#endif +void* MoEHighThroughputRuntime::resolveRecvXBuffer(int numTokens, int numRecvTokens, int hidden, + int xElementSize) const { + if (not available_ or not canUseDirectRecvPool(numTokens, numRecvTokens, hidden, xElementSize)) return nullptr; + return static_cast(recvPoolPtrs_[rank_]) + Config::get_recv_pool_header_bytes(numRanks_); } -int MoEHighThroughputRuntime::intranodeNotifyDispatch(int* rankPrefixMatrix, int* channelPrefixMatrix, - int* numRecvTokensPerExpert, const int* numTokensPerRank, - const int* numTokensPerExpert, const bool* isTokenInRank, - int numTokens, int numExperts, int xElementSize, - int expertAlignment, const Config& config, cudaStream_t stream) { - // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. - EP_HOST_ASSERT(config.num_sms % 2 == 0); - int dispatch_num_sms = 0, num_channels = 0; +int MoEHighThroughputRuntime::notifyDispatch(int* rankPrefixMatrix, int* channelPrefixMatrix, + int* numRecvTokensPerExpert, const int* numTokensPerRank, + const int* numTokensPerExpert, const bool* isTokenInRank, int numTokens, + int numExperts, int xElementSize, int expertAlignment, + cudaStream_t stream) { + EP_HOST_ASSERT(available_); + EP_HOST_ASSERT(numExperts > 0 and numExperts % numRanks_ == 0); + const int num_local_experts = numExperts / numRanks_; + EP_HOST_ASSERT(num_local_experts <= NUM_MAX_LOCAL_EXPERTS); + + int dispatch_num_sms = 0; + int num_channels = 0; bool all_sender = false; - computeIntranodeChannels(xElementSize, config, dispatch_num_sms, all_sender, num_channels); + computeDispatchChannels(xElementSize, dispatch_num_sms, all_sender, num_channels); + const int num_memset_int = num_channels * numRanks_ * 4; - const int num_local_experts = numExperts / num_ranks; + *moeRecvCounter_ = -1; + for (int i = 0; i < num_local_experts; ++i) moeRecvExpertCounter_[i] = -1; + intranode::notify_dispatch(numTokensPerRank, moeRecvCounterMapped_, numRanks_, numTokensPerExpert, + moeRecvExpertCounterMapped_, numExperts, numTokens, isTokenInRank, channelPrefixMatrix, + rankPrefixMatrix, num_memset_int, expertAlignment, bufferPtrsGpu_, taskFifoPtrsGpu_, head_, + rank_, stream, num_channels); + moveFifoSlots(3); - // Make comm_stream wait for the caller's stream before launching. - stream_wait(comm_stream, stream); - - // Barrier or send sizes - // To clean: channel start/end offset, head and tail - int num_memset_int = num_channels * num_ranks * 4; - - // Send sizes - // Meta information: - // - Size prefix by ranks, shaped as `[num_ranks, num_ranks]` - // - Size prefix by experts (not used later), shaped as `[num_ranks, num_local_experts]` - // NOTES: no more token dropping in this version - *moe_recv_counter = -1; - for (int i = 0; i < num_local_experts; ++i) moe_recv_expert_counter[i] = -1; - EP_HOST_ASSERT(num_ranks * (num_ranks + num_local_experts) * sizeof(int) <= num_nvl_bytes); - intranode::notify_dispatch(numTokensPerRank, moe_recv_counter_mapped, num_ranks, numTokensPerExpert, - moe_recv_expert_counter_mapped, numExperts, numTokens, isTokenInRank, channelPrefixMatrix, - rankPrefixMatrix, num_memset_int, expertAlignment, buffer_ptrs_gpu, task_fifo_ptrs_gpu, - head, rank, comm_stream, num_channels); - move_fifo_slots(3); - - // Synchronize total received tokens and tokens per expert int num_recv_tokens = -1; - auto start_time = std::chrono::high_resolution_clock::now(); + const auto start = std::chrono::high_resolution_clock::now(); while (true) { - // Read total count - num_recv_tokens = static_cast(*moe_recv_counter); - - // Read per-expert count - bool ready = (num_recv_tokens >= 0); - for (int i = 0; i < num_local_experts and ready; ++i) ready &= moe_recv_expert_counter[i] >= 0; - + num_recv_tokens = static_cast(*moeRecvCounter_); + bool ready = num_recv_tokens >= 0; + for (int i = 0; i < num_local_experts and ready; ++i) ready &= moeRecvExpertCounter_[i] >= 0; if (ready) break; - - // Timeout check - if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time) - .count() > NUM_CPU_TIMEOUT_SECS) + if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start).count() > + NUM_CPU_TIMEOUT_SECS) throw std::runtime_error("DeepEP error: CPU recv timeout"); } - for (int i = 0; i < num_local_experts; ++i) numRecvTokensPerExpert[i] = moe_recv_expert_counter[i]; - - // Make the caller's stream wait for comm_stream so the prefix matrices are visible. - stream_wait(stream, comm_stream); + for (int i = 0; i < num_local_experts; ++i) numRecvTokensPerExpert[i] = moeRecvExpertCounter_[i]; return num_recv_tokens; } -void MoEHighThroughputRuntime::intranodeDispatch(void* recvX, float* recvXScales, int64_t* recvTopkIdx, - float* recvTopkWeights, int* recvSrcIdx, int* sendHead, - int* recvChannelPrefixMatrix, const void* x, const float* xScales, - const int64_t* topkIdx, const float* topkWeights, - const bool* isTokenInRank, const int* rankPrefixMatrix, - const int* channelPrefixMatrix, int numTokens, int hidden, int numTopk, - int numScales, int numExperts, int xElementSize, int numRecvTokens, - bool cachedMode, const Config& config, cudaStream_t stream) { - // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. - EP_HOST_ASSERT(config.num_sms % 2 == 0); - int dispatch_num_sms = 0, num_channels = 0; +void MoEHighThroughputRuntime::dispatch(void* recvX, float* recvXScales, int64_t* recvTopkIdx, float* recvTopkWeights, + int* recvSrcIdx, int* sendHead, int* recvChannelPrefixMatrix, const void* x, + const float* xScales, const int64_t* topkIdx, const float* topkWeights, + const bool* isTokenInRank, const int* rankPrefixMatrix, + const int* channelPrefixMatrix, int numTokens, int hidden, int numTopk, + int numScales, int numExperts, int xElementSize, int numRecvTokens, + bool cachedMode, cudaStream_t stream) { + EP_HOST_ASSERT(available_); + EP_HOST_ASSERT(hidden > 0 and xElementSize > 0); + EP_HOST_ASSERT(static_cast(hidden) * xElementSize <= maxHiddenBytes_); + EP_HOST_ASSERT((hidden * xElementSize) % sizeof(int4) == 0); + EP_HOST_ASSERT(numTopk >= 0 and numTopk <= Config::MaxTopk); + EP_HOST_ASSERT(numScales >= 0 and numScales <= Config::MaxScales); + + int dispatch_num_sms = 0; + int num_channels = 0; bool all_sender = false; - computeIntranodeChannels(xElementSize, config, dispatch_num_sms, all_sender, num_channels); - - // num_experts is 0 in cached mode (matches the original tail, where it was - // derived as `cached_mode ? 0 : num_tokens_per_expert->size(0)`). + computeDispatchChannels(xElementSize, dispatch_num_sms, all_sender, num_channels); const int num_experts = cachedMode ? 0 : numExperts; - - // Input optional pointers (nullptr == not present). - const int64_t* topk_idx_ptr = topkIdx; - const float* topk_weights_ptr = topkWeights; - const float* x_scales_ptr = xScales; - int num_topk = numTopk; - int num_scales = numScales; - - // Make comm_stream wait for the caller's stream before launching. - stream_wait(comm_stream, stream); - - // Barrier (cached mode only). The non-cached caller already ran - // intranodeNotifyDispatch (which issued notify_dispatch + the barrier). - int num_memset_int = num_channels * num_ranks * 4; + const int num_memset_int = num_channels * numRanks_ * 4; if (cachedMode) { - // Copy rank prefix matrix and clean flags - intranode::cached_notify_dispatch(rankPrefixMatrix, num_memset_int, buffer_ptrs_gpu, task_fifo_ptrs_gpu, head, rank, - num_ranks, comm_stream); - move_fifo_slots(2); + intranode::cached_notify_dispatch(rankPrefixMatrix, num_memset_int, bufferPtrsGpu_, taskFifoPtrsGpu_, head_, rank_, + numRanks_, stream); + moveFifoSlots(2); } - // Sender-direct (MSCCLPP_EP_INTRA_DIRECT): make recv_x a zero-copy view of this rank's - // peer-mapped recv pool so the sender writes hidden straight to its final slot, - // eliminating the 2-hop ring + receiver hidden drain. The caller resolves recvX (= the - // pool view) via resolveIntranodeRecvXBuffer when this gate is active. - void** ep_intra_recv_pool_ptrs = nullptr; -#ifdef EP_DISPATCH_NCCLEP - const size_t ep_intra_pool_header_bytes = config.get_recv_pool_header_bytes(num_ranks); - const char* e_intra_direct = std::getenv("MSCCLPP_EP_INTRA_DIRECT"); - const bool ep_intra_direct = e_intra_direct != nullptr && std::atoi(e_intra_direct) != 0 && - recv_pool_local_ptr_ != nullptr && recv_pool_ptrs_gpu != nullptr && - numRecvTokens <= Config::kEpRecvPoolMaxTokens && - static_cast(numRecvTokens) * hidden * static_cast(xElementSize) <= - static_cast(Config::recv_pool_bytes_static(num_ranks)) - - static_cast(ep_intra_pool_header_bytes); -#else - const size_t ep_intra_pool_header_bytes = 0; - const bool ep_intra_direct = false; -#endif - if (ep_intra_direct) { - ep_intra_recv_pool_ptrs = recv_pool_ptrs_gpu; - // recvX already points at recv_pool_local_ptr_ + header (caller-resolved). - } + const bool direct = canUseDirectRecvPool(numTokens, numRecvTokens, hidden, xElementSize); + directDispatchReady_ = direct; + void** direct_pool_ptrs = direct ? recvPoolPtrsGpu_ : nullptr; + const size_t pool_header_bytes = Config::get_recv_pool_header_bytes(numRanks_); + if (direct) EP_HOST_ASSERT(recvX == static_cast(recvPoolPtrs_[rank_]) + pool_header_bytes); - // Caller-provided recv output pointers (no torch allocation here). - void* recv_x = recvX; - int64_t* recv_topk_idx_ptr = recvTopkIdx; - float* recv_topk_weights_ptr = recvTopkWeights; - float* recv_x_scales_ptr = recvXScales; - - // Dispatch - const int dispatch_hidden_int4 = static_cast(hidden * xElementSize / sizeof(int4)); + const int hidden_int4 = hidden * xElementSize / sizeof(int4); if (all_sender) { - // All-sender direct path: every block sends hidden straight to the dest pools; the - // hidden ring + receiver are unused. Requires the direct pool (assert it resolved). - EP_HOST_ASSERT(ep_intra_direct and ep_intra_recv_pool_ptrs != nullptr); - // The TMA combine is token-parallel and ignores channel_prefix_matrix; zero the recv - // copy so any (non-TMA) consumer sees a defined tensor. - CUDA_CHECK(cudaMemsetAsync(recvChannelPrefixMatrix, 0, static_cast(num_ranks) * num_channels * sizeof(int), - comm_stream)); - const int64_t meta_base = static_cast(Config::get_recv_pool_meta_base(num_ranks)); - const int64_t meta_slot_bytes = Config::kEpRecvPoolMetaBytes; - intranode::dispatch_allsender(sendHead, x, topk_idx_ptr, topk_weights_ptr, x_scales_ptr, isTokenInRank, - channelPrefixMatrix, numTokens, dispatch_hidden_int4, num_topk, num_experts, - num_scales, buffer_ptrs_gpu, rank, num_ranks, comm_stream, num_channels, - recv_pool_ptrs_gpu, static_cast(ep_intra_pool_header_bytes), meta_base, - meta_slot_bytes, ep_combine_recv_idx_gpu); - // Unpack the per-token metadata the senders packed into our pool META region. - intranode::intranode_meta_drain(recv_pool_local_ptr_, meta_base, numRecvTokens, recvSrcIdx, recv_topk_idx_ptr, - recv_topk_weights_ptr, recv_x_scales_ptr, num_topk, num_scales, meta_slot_bytes, - comm_stream); - } else { - EP_HOST_ASSERT( - num_ranks * num_ranks * sizeof(int) + // Size prefix matrix - num_channels * num_ranks * sizeof(int) + // Channel start offset - num_channels * num_ranks * sizeof(int) + // Channel end offset - num_channels * num_ranks * sizeof(int) * 2 + // Queue head and tail - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * hidden * xElementSize + // Data buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(int) + // Source index buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * - sizeof(int64_t) + // Top-k index buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * - sizeof(float) + // Top-k weight buffer - num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * sizeof(float) * - num_scales // FP8 scale buffer - <= num_nvl_bytes); - intranode::dispatch(recv_x, recv_x_scales_ptr, recvSrcIdx, recv_topk_idx_ptr, recv_topk_weights_ptr, - recvChannelPrefixMatrix, sendHead, x, x_scales_ptr, topk_idx_ptr, topk_weights_ptr, - isTokenInRank, channelPrefixMatrix, numTokens, dispatch_hidden_int4, num_topk, num_experts, - num_scales, buffer_ptrs_gpu, rank, num_ranks, comm_stream, dispatch_num_sms, - config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, - ep_intra_recv_pool_ptrs, static_cast(ep_intra_pool_header_bytes), - ep_intra_direct ? ep_combine_recv_idx_gpu : nullptr); + EP_HOST_ASSERT(direct); + CUDA_CHECK(cudaMemsetAsync(recvChannelPrefixMatrix, 0, static_cast(numRanks_) * num_channels * sizeof(int), + stream)); + intranode::dispatch_allsender( + sendHead, x, topkIdx, topkWeights, xScales, isTokenInRank, channelPrefixMatrix, numTokens, hidden_int4, numTopk, + num_experts, numScales, bufferPtrsGpu_, rank_, numRanks_, stream, num_channels, recvPoolPtrsGpu_, + static_cast(pool_header_bytes), static_cast(Config::get_recv_pool_meta_base(numRanks_)), + Config::RecvPoolMetaBytes, combineRecvIdxGpu_); + intranode::intranode_meta_drain( + recvPoolPtrs_[rank_], static_cast(Config::get_recv_pool_meta_base(numRanks_)), numRecvTokens, + recvSrcIdx, recvTopkIdx, recvTopkWeights, recvXScales, numTopk, numScales, Config::RecvPoolMetaBytes, stream); + return; } - // Make the caller's stream wait for comm_stream so the recv outputs are visible. - stream_wait(stream, comm_stream); + const size_t required_bytes = + static_cast(numRanks_) * numRanks_ * sizeof(int) + + static_cast(num_channels) * numRanks_ * sizeof(int) * 4 + + static_cast(num_channels) * numRanks_ * config_.num_max_nvl_chunked_recv_tokens * + (static_cast(hidden) * xElementSize + sizeof(int) + static_cast(numTopk) * sizeof(int64_t) + + static_cast(numTopk) * sizeof(float) + static_cast(numScales) * sizeof(float)); + EP_HOST_ASSERT(required_bytes <= ringBufferBytes_); + intranode::dispatch(recvX, recvXScales, recvSrcIdx, recvTopkIdx, recvTopkWeights, recvChannelPrefixMatrix, sendHead, + x, xScales, topkIdx, topkWeights, isTokenInRank, channelPrefixMatrix, numTokens, hidden_int4, + numTopk, num_experts, numScales, bufferPtrsGpu_, rank_, numRanks_, stream, dispatch_num_sms, + config_.num_max_nvl_chunked_send_tokens, config_.num_max_nvl_chunked_recv_tokens, + direct_pool_ptrs, static_cast(pool_header_bytes), direct ? combineRecvIdxGpu_ : nullptr); } -void MoEHighThroughputRuntime::intranodeCombine(void* combinedX, float* combinedTopkWeights, const void* x, - const float* topkWeights, const int* srcIdx, - const int* rankPrefixMatrix, const int* channelPrefixMatrix, - const int* sendHead, int numTokens, int numRecvTokens, int hidden, - int numTopk, int xElementSize, int ringNumChannels, - const Config& config, cudaStream_t stream) { - // One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving. - EP_HOST_ASSERT(config.num_sms % 2 == 0); - int num_channels = config.num_sms / 2; - - // The 2-hop ring fallback consumes channel_prefix_matrix; the TMA direct-gather combine - // ignores it. The all-sender dispatch produces a wider matrix (one column per block), and - // DISPATCH_NSM may produce fewer, so the column count is validated only inside the ring - // fallback (where it must fit the ring buffers); the TMA path tolerates any column count. - const int ring_num_channels = ringNumChannels; +void MoEHighThroughputRuntime::combine(void* combinedX, float* combinedTopkWeights, const void* x, + const float* topkWeights, const int* srcIdx, const int* rankPrefixMatrix, + const int* channelPrefixMatrix, const int* sendHead, int numTokens, + int numRecvTokens, int hidden, int numTopk, int xElementSize, + int ringNumChannels, cudaStream_t stream) { + EP_HOST_ASSERT(available_); + EP_HOST_ASSERT(static_cast(hidden) * xElementSize <= maxHiddenBytes_); EP_HOST_ASSERT((hidden * xElementSize) % sizeof(int4) == 0); + EP_HOST_ASSERT(numTopk >= 0 and numTopk <= Config::MaxTopk); - // Make comm_stream wait for the caller's stream before launching. - stream_wait(comm_stream, stream); - - const float* topk_weights_ptr = topkWeights; - float* recv_topk_weights_ptr = combinedTopkWeights; - int num_topk = numTopk; - - // Combine output. Default: TMA-staged direct-gather (MSCCLPP_EP_COMBINE_TMA != 0) reads - // each token's contributions straight from the peer recv pools through a cp.async.bulk SMEM - // pipeline. The grid is token-parallel, so MSCCLPP_EP_COMBINE_NSM sets the block count - // independently of the dispatch channel count. Falls back to the 2-hop ring combine when the - // direct-gather inputs are unavailable (no INTRA_DIRECT recv pools) or the escape hatch is set. - void* recv_x = combinedX; bool used_tma_combine = false; - { - static const bool ep_combine_tma_disabled = [] { - const char* e = std::getenv("MSCCLPP_EP_COMBINE_TMA"); - return e != nullptr and std::atoi(e) == 0; - }(); - // HT combine operates on BF16 (xElementSize == 2 mirrors the original - // `x.scalar_type() == torch::kBFloat16` gate); the TMA combine requires it. - const bool tma_inputs_ready = - recv_pool_ptrs_gpu != nullptr and ep_combine_recv_idx_gpu != nullptr and xElementSize == 2; - if (not ep_combine_tma_disabled and tma_inputs_ready) { - int combine_sms = config.num_sms; - const char* e_cnsm = std::getenv("MSCCLPP_EP_COMBINE_NSM"); - if (e_cnsm != nullptr and std::atoi(e_cnsm) >= 1) combine_sms = std::atoi(e_cnsm); - const int64_t intra_pool_header_bytes = static_cast(config.get_recv_pool_header_bytes(num_ranks)); - used_tma_combine = intranode::combine_tma( - CUDA_R_16BF, recv_x, recv_topk_weights_ptr, const_cast(sendHead), numRecvTokens, hidden, num_topk, - num_ranks, recv_pool_ptrs_gpu, ep_combine_recv_idx_gpu, intra_pool_header_bytes, combine_sms, comm_stream); - } + static const bool TmaDisabled = [] { + const char* env = std::getenv("MSCCLPP_EP_COMBINE_TMA"); + return env != nullptr and std::atoi(env) == 0; + }(); + if (not TmaDisabled and directDispatchReady_ and xElementSize == 2) { + EP_HOST_ASSERT(numTokens <= Config::RecvPoolMaxTokens); + int combine_sms = config_.num_sms; + if (const char* env = std::getenv("MSCCLPP_EP_COMBINE_NSM")) + if (std::atoi(env) >= 1) combine_sms = std::atoi(env); + used_tma_combine = intranode::combine_tma( + CUDA_R_16BF, combinedX, combinedTopkWeights, const_cast(sendHead), numRecvTokens, hidden, numTopk, + numRanks_, recvPoolPtrsGpu_, combineRecvIdxGpu_, + static_cast(Config::get_recv_pool_header_bytes(numRanks_)), combine_sms, stream); } - if (not used_tma_combine) { - // 2-hop ring fallback. Uses the channel count dispatch actually produced (ring_num_channels). - // The ring buffers are sized for config.num_sms/2 channels, so the all-sender dispatch - // (which produces more columns and writes hidden directly, not into the ring) is - // incompatible with this fallback -- it requires the TMA combine. - EP_HOST_ASSERT(ring_num_channels <= num_channels); - const int ring_num_sms = ring_num_channels * 2; - EP_HOST_ASSERT(ring_num_channels * num_ranks * sizeof(int) * 2 <= num_nvl_bytes); - intranode::cached_notify_combine(buffer_ptrs_gpu, const_cast(sendHead), ring_num_channels, numRecvTokens, - ring_num_channels * num_ranks * 2, task_fifo_ptrs_gpu, head, rank, num_ranks, - comm_stream); + if (used_tma_combine) return; - // NOTES: this function uses two FIFO slots (barrier before and after) - move_fifo_slots(2); + const int max_ring_channels = config_.num_sms / 2; + EP_HOST_ASSERT(ringNumChannels > 0 and ringNumChannels <= max_ring_channels); + intranode::cached_notify_combine(bufferPtrsGpu_, const_cast(sendHead), ringNumChannels, numRecvTokens, + ringNumChannels * numRanks_ * 2, taskFifoPtrsGpu_, head_, rank_, numRanks_, stream); + moveFifoSlots(2); - EP_HOST_ASSERT(ring_num_channels * num_ranks * sizeof(int) * 2 + // Queue head and tail - ring_num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * hidden * - xElementSize + // Data buffer - ring_num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * - sizeof(int) + // Source index buffer - ring_num_channels * num_ranks * config.num_max_nvl_chunked_recv_tokens * num_topk * - sizeof(float) // Top-k weight buffer - <= num_nvl_bytes); - intranode::combine(CUDA_R_16BF, recv_x, recv_topk_weights_ptr, x, topk_weights_ptr, srcIdx, rankPrefixMatrix, - channelPrefixMatrix, const_cast(sendHead), numTokens, numRecvTokens, hidden, num_topk, - buffer_ptrs_gpu, rank, num_ranks, comm_stream, ring_num_sms, - config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens); - } - - // Make the caller's stream wait for comm_stream so combinedX is visible. - stream_wait(stream, comm_stream); -} - -// ----------------------------------------------------------------------------- -// Internode (NVLink + RDMA) high-throughput path. Ported from DeepEP -// `csrc/deep_ep.cpp`; the kernels it drives are in -// `src/ext/ep/ht/kernels/internode.cu`. Validated end-to-end on 2 x H100 x 8 -// via `test/python/ep/test_internode_multirank.py`. De-torched the same -// way as the intranode path: tensor params became raw pointers + size ints, -// output tensors became caller pointers, the EventHandle / async / record_stream -// machinery became comm_stream stream_wait brackets, and the original single -// `internode_dispatch` was split into internodeNotifyDispatch (the non-cached -// notify phase) + internodeDispatch (the dispatch-kernel tail), mirroring the -// intranode split (see SIGNATURES.md for the notify->dispatch boundary). The -// kernels, env knobs, `#ifdef EP_DISPATCH_NCCLEP` blocks, recv-pool / VMM / NVLS -// logic and diagnostics are preserved verbatim. -// ----------------------------------------------------------------------------- - -void* MoEHighThroughputRuntime::resolveInternodeRecvXBuffer(int numRecvTokens, int hidden, int xElementSize, - const Config& config) const { - // Increment 4 (VMM pool): in the non-cached internode dispatch, recv_x is a zero-copy view of - // this rank's peer-mapped VMM recv-output pool (recv_pool_local_ptr_ + header) so the cross-GPU - // forwarder writes hidden straight to the destination's recv_x via TMA-eligible peer VAs and the - // direct-gather combine reads it back. Mirrors the `ep_use_direct` gate in internodeDispatch - // exactly (minus the cached-mode term, since the helper is only consulted on the non-cached - // forward path); returns nullptr when the caller should allocate recvX itself. -#ifdef EP_DISPATCH_NCCLEP - const size_t ep_pool_header_bytes = config.get_recv_pool_header_bytes(num_ranks); - const bool ep_use_direct = num_nvl_bytes > 0 and recv_pool_local_ptr_ != nullptr and recv_pool_ptrs_gpu != nullptr and - numRecvTokens <= Config::kEpRecvPoolMaxTokens; - (void)hidden; - (void)xElementSize; - if (ep_use_direct) { - return static_cast(recv_pool_local_ptr_) + ep_pool_header_bytes; - } - return nullptr; -#else - (void)numRecvTokens; - (void)hidden; - (void)xElementSize; - (void)config; - return nullptr; -#endif -} - -int MoEHighThroughputRuntime::getInternodeDispatchNumChannels(const Config& config) const { -#ifdef EP_DISPATCH_NCCLEP - return ep_flat_dispatch_channels(config.num_sms); -#else - return config.num_sms / 2; -#endif -} - -int MoEHighThroughputRuntime::getSourceMetaBytes() const { return internode::get_source_meta_bytes(); } - -int MoEHighThroughputRuntime::getNumMaxNvlPeers() const { return NUM_MAX_NVL_PEERS; } - -int MoEHighThroughputRuntime::internodeNotifyDispatch( - int* rdmaChannelPrefixMatrix, int* recvRdmaRankPrefixSum, int* gblChannelPrefixMatrix, int* recvGblRankPrefixSum, - int* numRecvTokensPerExpert, int* numRdmaRecvTokens, const int* numTokensPerRank, const int* numTokensPerRdmaRank, - const int* numTokensPerExpert, const bool* isTokenInRank, int numTokens, int numExperts, int hidden, int numScales, - int numTopk, int xElementSize, int expertAlignment, const Config& config, cudaStream_t stream) { - // In dispatch, CPU will busy-wait until GPU receive tensor size metadata from other ranks, which can be quite long. - int num_channels = config.num_sms / 2; - EP_HOST_ASSERT(config.num_sms % 2 == 0); - EP_HOST_ASSERT(0 < num_rdma_ranks and num_rdma_ranks <= NUM_MAX_RDMA_PEERS); - - // inc6/inc7 (kEpFlat): MSCCLPP_EP_DISPATCH_NSM sets the flat dispatch block count - // (== num_channels, since the flat all-sender path launches one sender block per - // channel) independently of config.num_sms. num_channels is the token-partitioning - // granularity shared by notify_dispatch, the prefix-matrix allocations, and the - // dispatch grid, so resolving it here keeps the whole dispatch pipeline self-consistent - // (the NSM sweep already varied num_channels 8..76 with byte-correct recv). Clamped to - // [1, config.num_sms] (the flat path has no forwarder, so it can use the FULL SM budget: - // num_sms=16 + MSCCLPP_EP_DISPATCH_NSM=16 -> 16 blocks, matching NCCL-EP); the matching - // RDMA/NVL buffer is sized for the same channel count via ep_buffer_channels(). Unset -> - // num_sms/2 (byte-identical default). Flat-only; the 2-hop path keeps num_sms/2. The notify - // phase is always non-cached, so the original `if (not cached_mode)` guard is unconditional here. -#ifdef EP_DISPATCH_NCCLEP - num_channels = ep_flat_dispatch_channels(config.num_sms); -#endif - - const int num_local_experts = numExperts / num_ranks; - const int hidden_int4 = static_cast(hidden * xElementSize / sizeof(int4)); - - // Make comm_stream wait for the caller's stream before launching. - stream_wait(comm_stream, stream); - - // Send sizes - *moe_recv_counter = -1, *moe_recv_rdma_counter = -1; - for (int i = 0; i < num_local_experts; ++i) moe_recv_expert_counter[i] = -1; - // NVLS Phase 2: bump the per-call epoch counter so the kernel's - // barrier spin uses a fresh expected value (epoch * num_ranks). - if (nvls_ht_enabled) ++nvls_ht_epoch; - internode::notify_dispatch( - numTokensPerRank, moe_recv_counter_mapped, num_ranks, numTokensPerRdmaRank, moe_recv_rdma_counter_mapped, - numTokensPerExpert, moe_recv_expert_counter_mapped, numExperts, isTokenInRank, numTokens, num_channels, - hidden_int4, numScales, numTopk, expertAlignment, rdmaChannelPrefixMatrix, recvRdmaRankPrefixSum, - gblChannelPrefixMatrix, recvGblRankPrefixSum, rdma_buffer_ptr, config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, config.num_max_nvl_chunked_recv_tokens, task_fifo_ptrs_gpu, head, rank, comm_stream, - config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), num_nvl_bytes, low_latency_mode, - port_channel_handles_device_ptr.get(), memory_channel_handles_device_ptr.get(), - nvls_ht_enabled ? nvls_ht_mc_ptr : nullptr, nvls_ht_enabled ? nvls_ht_dev_ptr : nullptr, nvls_ht_off_barrier, - nvls_ht_off_data, nvls_ht_epoch, kNvlsPerPeerBytes); - move_fifo_slots(3); - - // Synchronize total received tokens and tokens per expert - int num_recv_tokens = -1, num_rdma_recv_tokens = -1; - auto start_time = std::chrono::high_resolution_clock::now(); - while (true) { - num_recv_tokens = static_cast(*moe_recv_counter); - num_rdma_recv_tokens = static_cast(*moe_recv_rdma_counter); - - bool ready = (num_recv_tokens >= 0) and (num_rdma_recv_tokens >= 0); - for (int i = 0; i < num_local_experts and ready; ++i) ready &= moe_recv_expert_counter[i] >= 0; - - if (ready) break; - - if (std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time) - .count() > NUM_CPU_TIMEOUT_SECS) { - printf("Global rank: %d, num_recv_tokens: %d, num_rdma_recv_tokens: %d\n", rank, num_recv_tokens, - num_rdma_recv_tokens); - for (int i = 0; i < num_local_experts; ++i) - printf("moe_recv_expert_counter[%d]: %d\n", i, moe_recv_expert_counter[i]); - throw std::runtime_error("mscclpp::ep error: timeout (internode_dispatch CPU)"); - } - } - for (int i = 0; i < num_local_experts; ++i) numRecvTokensPerExpert[i] = moe_recv_expert_counter[i]; - *numRdmaRecvTokens = num_rdma_recv_tokens; - - // Make the caller's stream wait for comm_stream so the prefix matrices are visible. - stream_wait(stream, comm_stream); - return num_recv_tokens; -} - -void MoEHighThroughputRuntime::internodeDispatch( - void* recvX, float* recvXScales, int64_t* recvTopkIdx, float* recvTopkWeights, void* recvSrcMeta, - int* recvRdmaChannelPrefixMatrix, int* recvGblChannelPrefixMatrix, int* sendRdmaHead, int* sendNvlHead, - const void* x, const float* xScales, const int64_t* topkIdx, const float* topkWeights, const bool* isTokenInRank, - const int* rdmaChannelPrefixMatrix, const int* recvRdmaRankPrefixSum, const int* gblChannelPrefixMatrix, - const int* recvGblRankPrefixSum, int numTokens, int hidden, int numTopk, int numScales, int numExperts, - int xElementSize, int numRecvTokens, int numRdmaRecvTokens, bool cachedMode, const Config& config, - cudaStream_t stream) { - int num_channels = config.num_sms / 2; - EP_HOST_ASSERT(config.num_sms % 2 == 0); - EP_HOST_ASSERT(0 < num_rdma_ranks and num_rdma_ranks <= NUM_MAX_RDMA_PEERS); - - // inc6/inc7 (kEpFlat): MSCCLPP_EP_DISPATCH_NSM sets the flat dispatch block count - // (== num_channels, since the flat all-sender path launches one sender block per - // channel) independently of config.num_sms. num_channels is the token-partitioning - // granularity shared by notify_dispatch, the prefix-matrix allocations, and the - // dispatch grid, so resolving it here keeps the whole dispatch pipeline self-consistent - // (the NSM sweep already varied num_channels 8..76 with byte-correct recv). Clamped to - // [1, config.num_sms] (the flat path has no forwarder, so it can use the FULL SM budget: - // num_sms=16 + MSCCLPP_EP_DISPATCH_NSM=16 -> 16 blocks, matching NCCL-EP); the matching - // RDMA/NVL buffer is sized for the same channel count via ep_buffer_channels(). Unset -> - // num_sms/2 (byte-identical default). Flat-only; the 2-hop path keeps num_sms/2. -#ifdef EP_DISPATCH_NCCLEP - if (not cachedMode) num_channels = ep_flat_dispatch_channels(config.num_sms); -#endif - - const int hidden_int4 = static_cast(hidden * xElementSize / sizeof(int4)); - // num_experts is 0 in cached mode (matches the original tail, where it was - // derived as `cached_mode ? 0 : num_tokens_per_expert->size(0)`). - const int num_experts = cachedMode ? 0 : numExperts; - - // Input optional pointers (nullptr == not present). - const int64_t* topk_idx_ptr = topkIdx; - const float* topk_weights_ptr = topkWeights; - const float* x_scales_ptr = xScales; - int num_topk = numTopk; - int num_scales = numScales; - - // Make comm_stream wait for the caller's stream before launching. - stream_wait(comm_stream, stream); - - // Barrier or send sizes. The non-cached caller already ran internodeNotifyDispatch - // (which issued notify_dispatch + the host busy-wait); cached mode just needs the barrier. - if (cachedMode) { - // Just a barrier and clean flags - if (nvls_ht_enabled) ++nvls_ht_cached_epoch; - internode::cached_notify(hidden_int4, num_scales, num_topk, num_topk, num_ranks, num_channels, 0, nullptr, nullptr, - nullptr, nullptr, rdma_buffer_ptr, config.num_max_rdma_chunked_recv_tokens, - buffer_ptrs_gpu, config.num_max_nvl_chunked_recv_tokens, task_fifo_ptrs_gpu, head, rank, - comm_stream, config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), - num_nvl_bytes, true, low_latency_mode, port_channel_handles_device_ptr.get(), - memory_channel_handles_device_ptr.get(), nvls_ht_enabled ? nvls_ht_mc_ptr : nullptr, - nvls_ht_enabled ? nvls_ht_dev_ptr : nullptr, nvls_ht_off_barrier, nvls_ht_cached_epoch); - move_fifo_slots(2); - } - - // Allocate new tensors - void** ep_recv_pool_ptrs = nullptr; // non-null selects the increment-4 VMM direct-write path - void** ep_recv_pool_global_ptrs = nullptr; // inc5: domain-wide pool bases (sender direct-write) -#ifdef EP_DISPATCH_NCCLEP - // Increment 4 (VMM pool): when num_recv_tokens fits the fixed pool, back recv_x - // by the local VMM recv-output pool so the cross-GPU forwarder can write hidden - // straight to the destination's recv_x via TMA-eligible peer VAs. Publish our - // recv_gbl_rank_prefix_sum into the peer-readable pool header. recv_x (= recvX) is the - // caller-resolved pool view (recv_pool_local_ptr_ + header) from resolveInternodeRecvXBuffer. - const bool ep_use_direct = (not cachedMode) and num_nvl_bytes > 0 and recv_pool_local_ptr_ != nullptr and - recv_pool_ptrs_gpu != nullptr and numRecvTokens <= Config::kEpRecvPoolMaxTokens; - if (ep_use_direct) { - ep_recv_pool_ptrs = recv_pool_ptrs_gpu; - ep_recv_pool_global_ptrs = recv_pool_global_ptrs_gpu; // inc5: null unless MSCCLPP_EP_DIRECT exchanged - void* pool_base = recv_pool_local_ptr_; - CUDA_CHECK(cudaMemcpyAsync(pool_base, recvGblRankPrefixSum, static_cast(num_ranks) * sizeof(int), - cudaMemcpyDeviceToDevice, comm_stream)); - } -#endif - - // Caller-provided recv output pointers (no torch allocation here). The recv-side metadata - // outputs (recvSrcMeta / recv*ChannelPrefixMatrix / send*Head) are passed straight into the - // kernel with the original `cached_mode ? nullptr : ...` gating below. - void* recv_x = recvX; - int64_t* recv_topk_idx_ptr = recvTopkIdx; - float* recv_topk_weights_ptr = recvTopkWeights; - float* recv_x_scales_ptr = recvXScales; - - // Launch data dispatch - // Phase 3: pass NVLS counter region pointers (head/tail × mc/dev). When - // `nvls_ht_enabled` is false, all four are nullptr and the kernel falls - // back to the legacy PortChannel/atomicAdd path. - void* nvls_head_mc = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_mc_ptr) + nvls_ht_off_head) : nullptr; - void* nvls_head_dev = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_dev_ptr) + nvls_ht_off_head) : nullptr; - void* nvls_tail_mc = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_mc_ptr) + nvls_ht_off_tail) : nullptr; - void* nvls_tail_dev = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_dev_ptr) + nvls_ht_off_tail) : nullptr; - internode::dispatch(recv_x, recv_x_scales_ptr, recv_topk_idx_ptr, recv_topk_weights_ptr, - cachedMode ? nullptr : recvSrcMeta, x, x_scales_ptr, topk_idx_ptr, topk_weights_ptr, - cachedMode ? nullptr : sendRdmaHead, cachedMode ? nullptr : sendNvlHead, - cachedMode ? nullptr : recvRdmaChannelPrefixMatrix, - cachedMode ? nullptr : recvGblChannelPrefixMatrix, rdmaChannelPrefixMatrix, recvRdmaRankPrefixSum, - gblChannelPrefixMatrix, recvGblRankPrefixSum, numTokens, hidden_int4, num_scales, num_topk, - num_experts, isTokenInRank, rdma_buffer_ptr, config.num_max_rdma_chunked_send_tokens, - config.num_max_rdma_chunked_recv_tokens, buffer_ptrs_gpu, config.num_max_nvl_chunked_send_tokens, - config.num_max_nvl_chunked_recv_tokens, rank, num_ranks, cachedMode, comm_stream, num_channels, - low_latency_mode, port_channel_handles_device_ptr.get(), memory_channel_handles_device_ptr.get(), - nvls_head_mc, nvls_head_dev, nvls_tail_mc, nvls_tail_dev, peer_rdma_bases_gpu, ep_recv_pool_ptrs, - ep_recv_pool_global_ptrs, ep_recv_pool_global_ptrs ? ep_combine_recv_idx_gpu : nullptr); - -#ifdef EP_DISPATCH_NCCLEP - // Increment 6 (kEpFlat): when MSCCLPP_EP_FLAT is set, the sender wrote each token's - // metadata straight into the destination pool's META region. Drain it into the - // recv_* output tensors on comm_stream right after dispatch. With the receiver role - // still present (approach 1), the dispatch kernel does not exit until all senders' - // pool writes are visible, so this same-stream drain sees complete metadata (the - // same guarantee inc5's recv_x relies on). Overwrites the receiver's metadata output - // with identical values; this validates the flat producer before the receiver is - // removed in the all-sender stage. - { - const char* e_direct = std::getenv("MSCCLPP_EP_DIRECT"); - const char* e_flat = std::getenv("MSCCLPP_EP_FLAT"); - const bool ep_flat = (e_flat && std::atoi(e_flat) != 0) && (e_direct && std::atoi(e_direct) != 0); - if (ep_flat and ep_use_direct and not cachedMode and topk_idx_ptr != nullptr) { - const int64_t ep_meta_base = static_cast(Config::get_recv_pool_meta_base(num_ranks)); - internode::flat_meta_drain(recv_pool_local_ptr_, ep_meta_base, numRecvTokens, recvSrcMeta, recv_x_scales_ptr, - recv_topk_idx_ptr, recv_topk_weights_ptr, num_scales, num_topk, num_experts, num_ranks, - rank, Config::kEpRecvPoolMetaBytes, comm_stream); - } - } -#endif - - // Make the caller's stream wait for comm_stream so the recv outputs are visible. - stream_wait(stream, comm_stream); -} - -void MoEHighThroughputRuntime::internodeCombine(void* combinedX, float* combinedTopkWeights, const void* x, - const float* topkWeights, const void* srcMeta, - const bool* isCombinedTokenInRank, const int* rdmaChannelPrefixMatrix, - const int* rdmaRankPrefixSum, const int* gblChannelPrefixMatrix, - const int* combinedRdmaHead, const int* combinedNvlHead, int numTokens, - int numCombinedTokens, int hidden, int numTopk, int xElementSize, - const Config& config, cudaStream_t stream) { - const int num_channels = config.num_sms / 2; - EP_HOST_ASSERT(config.num_sms % 2 == 0); - - // inc6 (kEpFlat): detect the flat direct-gather combine up front. Under flat the - // gather path ignores the per-channel prefix matrices entirely (they are vestigial - // and, with MSCCLPP_EP_DISPATCH_NSM, may be sized for a different channel count), so - // their num_channels shape checks below are relaxed. The flag is reused later to skip - // cached_notify and to honor MSCCLPP_EP_COMBINE_NSM. - bool ep_flat_combine = false; -#ifdef EP_DISPATCH_NCCLEP - { - const char* e_direct = std::getenv("MSCCLPP_EP_DIRECT"); - const char* e_flat = std::getenv("MSCCLPP_EP_FLAT"); - const bool ep_flat = (e_flat && std::atoi(e_flat) != 0) && (e_direct && std::atoi(e_direct) != 0); - ep_flat_combine = ep_flat and recv_pool_global_ptrs_gpu != nullptr and ep_combine_recv_idx_gpu != nullptr; - } -#endif - - const int hidden_int4 = static_cast(hidden * xElementSize / sizeof(int4)); - EP_HOST_ASSERT((hidden * xElementSize) % sizeof(int4) == 0); - - // Make comm_stream wait for the caller's stream before launching. - stream_wait(comm_stream, stream); - - const float* topk_weights_ptr = topkWeights; - float* combined_topk_weights_ptr = combinedTopkWeights; - int num_topk = numTopk; - - EP_HOST_ASSERT(config.num_max_nvl_chunked_recv_tokens % num_rdma_ranks == 0); - EP_HOST_ASSERT(config.num_max_nvl_chunked_send_tokens <= config.num_max_nvl_chunked_recv_tokens / num_rdma_ranks); - - // inc6 (kEpFlat): under the flat all-sender path the combine kernel early-returns - // to the pure direct-gather branch (recv_pool_global_ptrs + ep_combine_recv_idx), - // so it never touches the rdma/nvl ring or the head breadcrumbs. cached_notify's - // ring cleanup + head-prep is therefore unnecessary AND unsafe: its sm_id>=3 branch - // indexes rdma_channel_prefix_matrix (the forwarder-produced recv_rdma_channel_prefix_matrix) - // to derive token ranges, then writes combined_nvl_head over that range. Under - // all-sender flat the forwarder is removed so those tensors are never written -> - // garbage ranges -> out-of-bounds writes (illegal memory access). Skip the whole - // cached_notify + fifo-advance step under flat; the direct gather needs none of it. - // (ep_flat_combine was computed once near the top of this function.) - if (not ep_flat_combine) { - internode::cached_notify( - hidden_int4, 0, 0, num_topk, num_ranks, num_channels, numCombinedTokens, const_cast(combinedRdmaHead), - rdmaChannelPrefixMatrix, rdmaRankPrefixSum, const_cast(combinedNvlHead), rdma_buffer_ptr, - config.num_max_rdma_chunked_recv_tokens, buffer_ptrs_gpu, config.num_max_nvl_chunked_recv_tokens, - task_fifo_ptrs_gpu, head, rank, comm_stream, - config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks), num_nvl_bytes, false, low_latency_mode, - port_channel_handles_device_ptr.get(), memory_channel_handles_device_ptr.get(), - nvls_ht_enabled ? nvls_ht_mc_ptr : nullptr, nvls_ht_enabled ? nvls_ht_dev_ptr : nullptr, nvls_ht_off_barrier, - (nvls_ht_enabled ? ++nvls_ht_cached_epoch : 0)); - move_fifo_slots(2); - } - - // inc6 (kEpFlat): the flat combine is a pure direct-gather whose grid is - // independent of the dispatch channel partitioning -- it strides every warp over - // num_combined_tokens and never indexes the per-channel prefix matrices. The - // combine SM sweep showed it saturates the NVLink ceiling at ~76 blocks, so the - // grid can be capped independently of config.num_sms via MSCCLPP_EP_COMBINE_NSM - // (combine block count; default keeps the dispatch-tied num_channels*2). Only - // applied under flat; the 2-hop combine MUST keep num_channels matching the - // prefix-matrix column count, so its grid is left untouched. Capped at - // num_channels (=config.num_sms/2) so the combine grid never exceeds the SM budget. - int combine_grid_channels = num_channels; -#ifdef EP_DISPATCH_NCCLEP - if (ep_flat_combine) { - const char* e_cnsm = std::getenv("MSCCLPP_EP_COMBINE_NSM"); - if (e_cnsm != nullptr) { - int cnsm_blocks = std::atoi(e_cnsm); - if (cnsm_blocks >= 2) { - int cnsm_channels = cnsm_blocks / 2; - combine_grid_channels = (cnsm_channels < num_channels) ? cnsm_channels : num_channels; - } - } - } -#endif - - void* combined_x = combinedX; - // Phase 3: NVLS counter region pointers for combine kernel. - void* combine_nvls_head_mc = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_mc_ptr) + nvls_ht_off_head) : nullptr; - void* combine_nvls_head_dev = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_dev_ptr) + nvls_ht_off_head) : nullptr; - void* combine_nvls_tail_mc = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_mc_ptr) + nvls_ht_off_tail) : nullptr; - void* combine_nvls_tail_dev = - nvls_ht_enabled ? static_cast(static_cast(nvls_ht_dev_ptr) + nvls_ht_off_tail) : nullptr; - internode::combine(CUDA_R_16BF, combined_x, combined_topk_weights_ptr, isCombinedTokenInRank, x, topk_weights_ptr, - combinedRdmaHead, combinedNvlHead, srcMeta, rdmaChannelPrefixMatrix, rdmaRankPrefixSum, - gblChannelPrefixMatrix, numTokens, numCombinedTokens, hidden, num_topk, rdma_buffer_ptr, - config.num_max_rdma_chunked_send_tokens, config.num_max_rdma_chunked_recv_tokens, buffer_ptrs_gpu, - config.num_max_nvl_chunked_send_tokens, config.num_max_nvl_chunked_recv_tokens, rank, num_ranks, - comm_stream, combine_grid_channels, low_latency_mode, port_channel_handles_device_ptr.get(), - memory_channel_handles_device_ptr.get(), combine_nvls_head_mc, combine_nvls_head_dev, - combine_nvls_tail_mc, combine_nvls_tail_dev, peer_rdma_bases_gpu, recv_pool_global_ptrs_gpu, - ep_combine_recv_idx_gpu); - - // Make the caller's stream wait for comm_stream so combinedX is visible. - stream_wait(stream, comm_stream); + const size_t required_bytes = + static_cast(ringNumChannels) * numRanks_ * sizeof(int) * 2 + + static_cast(ringNumChannels) * numRanks_ * config_.num_max_nvl_chunked_recv_tokens * + (static_cast(hidden) * xElementSize + sizeof(int) + static_cast(numTopk) * sizeof(float)); + EP_HOST_ASSERT(required_bytes <= ringBufferBytes_); + intranode::combine(CUDA_R_16BF, combinedX, combinedTopkWeights, x, topkWeights, srcIdx, rankPrefixMatrix, + channelPrefixMatrix, const_cast(sendHead), numTokens, numRecvTokens, hidden, numTopk, + bufferPtrsGpu_, rank_, numRanks_, stream, ringNumChannels * 2, + config_.num_max_nvl_chunked_send_tokens, config_.num_max_nvl_chunked_recv_tokens); } } // namespace ep diff --git a/src/ext/ep/ht_runtime.hpp b/src/ext/ep/ht_runtime.hpp index 4066f5d3..15a16ecc 100644 --- a/src/ext/ep/ht_runtime.hpp +++ b/src/ext/ep/ht_runtime.hpp @@ -3,381 +3,87 @@ // // Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP) // branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. -// -// High-throughput (HT) MoE dispatch/combine runtime. This is a torch-free, -// raw-pointer port of the DeepEP-style `Buffer` (formerly `src/ext/ep/ht/ -// buffer.{cc,hpp}`), following the same boundary convention as the -// low-latency `MoERuntime` (`src/ext/ep/moe_runtime.{cc,hpp}`): tensor params -// become raw device pointers + explicit size ints, output tensors become -// caller-provided pointers, and torch CUDA streams become `cudaStream_t`. The -// CUDA kernels it drives (`intranode::` / `internode::`) are unchanged. #pragma once #include +#include #include -#include #include -#include -#include -#include -#include -#include -#include #include -#include "constants.cuh" -#include "exception.cuh" #include "ht/config.hpp" namespace mscclpp { namespace ep { -// High-throughput expert-parallel runtime. Sibling to `MoERuntime` (low -// latency); both expose a raw-pointer, torch-free boundary so the extension -// module never links libtorch. -// -// Dynamic recv sizing is handled with an explicit two-phase API: the caller -// first runs `intranodeNotifyDispatch` (returns `numRecvTokens`), allocates the -// recv output buffers, then runs `intranodeDispatch`. The cached fast path -// skips the notify phase (`cachedMode == true`). -struct MoEHighThroughputRuntime { - EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8 || NUM_MAX_NVL_PEERS == 4, - "The number of maximum NVLink peers must be 4 or 8"); - - private: - // Low-latency mode buffer - int low_latency_buffer_idx = 0; - // Retained as a compile-time-false member so the original construction / - // sync / dispatch bodies (which branch on `low_latency_mode`) port verbatim. - // The HT runtime is always high-throughput; the low-latency path is served - // by `MoERuntime`. - bool low_latency_mode = false; - - // NVLink Buffer - int64_t num_nvl_bytes; - void* buffer_ptrs[NUM_MAX_NVL_PEERS] = {nullptr}; - void** buffer_ptrs_gpu = nullptr; - // Increment 3: byte offset of the peer-mapped recv-output pool within the NVL - // allocation (buffer_ptrs[*] + recv_pool_off_). -1 until set in the ctor. - int64_t recv_pool_off_ = -1; - // Increment 4: VMM-allocated (cuMem FABRIC/POSIX-FD via gpuCallocPhysical) - // recv-output pool. recv_pool_local_ptr_ is this rank's local pool; the peer - // bases (imported via registerMemory/recvMemory) live in recv_pool_ptrs_ / - // recv_pool_ptrs_gpu. These are TMA-eligible peer VAs (unlike cudaIpc maps). - // Non-null recv_pool_local_ptr_ selects the increment-4 VMM direct-write path. - void* recv_pool_local_ptr_ = nullptr; - std::vector recv_pool_ptrs_; - void** recv_pool_ptrs_gpu = nullptr; - // Keep imported peer RegisteredMemory alive so the cuMem mapping persists for - // the runtime's lifetime (recv_pool_ptrs_[*] alias their .data()). - std::vector recv_pool_remote_mems_; - // Increment 5 (inc5 flat-domain dispatch): domain-wide recv-pool bases indexed - // by GLOBAL rank (all num_ranks across the NVLink domain), so the RDMA sender - // can write each token directly into the destination GPU's recv pool. Populated - // only when MSCCLPP_EP_DIRECT is set. recv_pool_global_ptrs_[rank]==local pool. - std::vector recv_pool_global_ptrs_; - void** recv_pool_global_ptrs_gpu = nullptr; - std::vector recv_pool_global_remote_mems_; - // Increment 5 combine-direct (Stage 1): per-(source token, dst global rank) - // recv-pool slot index written by the dispatch sender (= ep_my_idx). Lets the - // combine path gather each token's contributions straight from the peer pools. - // Allocated [kEpRecvPoolMaxTokens * num_ranks] ints under MSCCLPP_EP_DIRECT. - int* ep_combine_recv_idx_gpu = nullptr; - - // NVSHMEM Buffer - int64_t num_rdma_bytes; - void* rdma_buffer_ptr = nullptr; - - // Device info and communication - int device_id; - int rank, rdma_rank, nvl_rank; - int num_ranks, num_rdma_ranks, num_nvl_ranks; - cudaIpcMemHandle_t ipc_handles[NUM_MAX_NVL_PEERS]; - - // Stream for communication. Replaces the torch `at::cuda::CUDAStream` - // member; created in the ctor (`cudaStreamCreateWithFlags`) and destroyed in - // the dtor. The kernels still launch on this stream exactly as before. - cudaStream_t comm_stream = nullptr; - // Reusable event for cross-stream ordering between the caller-provided stream - // and `comm_stream` (replaces the torch EventHandle / stream_wait dance). - cudaEvent_t comm_event = nullptr; - - // After IPC/NVSHMEM synchronization, this flag will be true - bool available = false; - - // Task fifo - int head = 0; - int* task_fifo_ptrs[NUM_MAX_NVL_PEERS] = {nullptr}; - int** task_fifo_ptrs_gpu = nullptr; - - // Workspace - void* workspace = nullptr; - - // Host-side MoE info - volatile int* moe_recv_counter = nullptr; - int* moe_recv_counter_mapped = nullptr; - - // Host-side expert-level MoE info - volatile int* moe_recv_expert_counter = nullptr; - int* moe_recv_expert_counter_mapped = nullptr; - - // Host-side RDMA-level MoE info - volatile int* moe_recv_rdma_counter = nullptr; - int* moe_recv_rdma_counter_mapped = nullptr; - - std::shared_ptr bootstrap; - // One ProxyService spawns a single proxy thread that drains every PortChannel - // FIFO it owns. With LL combine pushing thousands of triggers per iter, the - // single thread becomes the wall-clock bottleneck on cross-node runs. We - // shard channels across `proxy_services` so each gets its own thread/FIFO, - // increasing host-side dispatch parallelism (no kernel changes required). - // Count is resolved at construction (env `MSCCLPP_EP_NUM_PROXIES` or - // arch-aware default). - int num_proxy_services = 1; - std::vector> proxy_services; - std::shared_ptr communicator; - std::vector port_channels; - std::vector memory_channels; - std::shared_ptr port_channel_handles_device_ptr; - std::shared_ptr memory_channel_handles_device_ptr; - - // LL fast path: peer-mapped RDMA buffer pointers. - // ``peer_rdma_bases[r]`` aliases rank ``r``'s ``rdma_buffer_ptr`` through - // mscclpp's CudaIpc transport. Intranode peers use POSIX-FD CUDA IPC; - // cross-node peers use cuMem fabric handles routed through nvidia-imex - // over the NVL72 NVSwitch fabric (Proposal A — replaces RDMA atomicAdd - // with NVLink atomics, since Azure CX-7 RoCE has IBV_ATOMIC_NONE). - // Populated in ``sync()`` when ``low_latency_mode``; empty otherwise. - std::vector peer_rdma_bases; - void** peer_rdma_bases_gpu = nullptr; - // Base MemoryChannels over CUDA IPC used only for the LL barrier ring. - std::vector ll_memory_channels; - std::shared_ptr ll_memory_channel_handles_device_ptr; - int ll_ranks_per_ipc_domain = 0; - bool ll_ipc_ready = false; - - // NVLS multicast for HT internode (Wide Proposal B2). - // - // When `mscclpp::isNvlsSupported()` is true and `num_rdma_ranks > 1`, - // we set up a multicast-bound buffer carrying: - // - tail counters[num_channels][num_rdma_ranks][num_rdma_ranks] uint64_t - // - head counters[num_channels][num_rdma_ranks][num_rdma_ranks] uint64_t - // - notify_dispatch barrier epoch[num_rdma_ranks] uint64_t - // - notify_dispatch small-data slots[num_rdma_ranks][kSummaryBytes] - // - // Cross-node atomic adds use `multimem.red.add.u64` PTX which travels - // over the NVL72 fabric instead of broken IB atomics on Azure CX-7 RoCE. - // The kernels select between this NVLS path and the legacy PortChannel - // path at runtime based on `nvls_ht_enabled`. - // - // Falls back gracefully on platforms without NVLS multicast support - // (e.g. H100+IB, A100+IB clusters): `nvls_ht_enabled` stays `false`, - // all NVLS pointers stay `nullptr`, and the original PortChannel - // signal/wait + atomicAdd path remains active. - bool nvls_ht_enabled = false; - std::shared_ptr nvls_ht_conn; - // SwitchChannel keeps the multicast pointer alive (its destructor - // unbinds the multicast); device pointers below are extracted from it. - std::shared_ptr nvls_ht_sc; - // Underlying GpuBuffer (multicast-eligible physical alloc); kept alive - // for the lifetime of the multicast binding. - std::shared_ptr> nvls_ht_buffer; - // mc_ptr: multicast-side device pointer (writes hit all peers via switch). - // dev_ptr: local-side device pointer (reads see local copy of the same - // physical memory). - void* nvls_ht_mc_ptr = nullptr; - void* nvls_ht_dev_ptr = nullptr; - // Sub-region byte offsets within the multicast buffer (set in sync()). - size_t nvls_ht_off_tail = 0; - size_t nvls_ht_off_head = 0; - size_t nvls_ht_off_barrier = 0; - size_t nvls_ht_off_data = 0; - size_t nvls_ht_total_bytes = 0; - // Per-call epoch counter for NVLS barrier slots. Incremented on the host - // before each kernel launch that uses an NVLS barrier; the kernel spins - // until the barrier slot reaches `epoch * num_ranks`. - uint64_t nvls_ht_epoch = 0; - // Independent epoch for cached_notify barrier slots (offsets +24 / +32), - // since those slots are only touched when the cached path is taken — using - // the shared `nvls_ht_epoch` would over-count the expected value relative - // to the number of times those particular slots have actually been bumped. - uint64_t nvls_ht_cached_epoch = 0; - // Worst-case shape parameters used to size the buffer: - // stride_per_channel = num_rdma_ranks * num_rdma_ranks (counter slots) - // We allocate for `kNvlsMaxChannels` so any `num_sms` config fits. - static constexpr int kNvlsMaxChannels = 64; // num_sms / 2 upper bound - static constexpr int kNvlsPerPeerBytes = 1024; // small-data per (sender, receiver) pair - // Number of distinct barrier slots in the barrier sub-region (each u64). - static constexpr int kNvlsBarrierSlots = 8; - - private: - void move_fifo_slots(int num_slots = 1); - - // Cross-stream ordering helper: makes `dst` wait for all work currently - // enqueued on `src`, using the reusable `comm_event`. Replaces the torch - // `stream_wait` / EventHandle machinery. - void stream_wait(cudaStream_t dst, cudaStream_t src); - - // Resolve the intranode dispatch channel layout from the config + env knobs - // (MSCCLPP_EP_DISPATCH_NSM / MSCCLPP_EP_INTRA_DIRECT / MSCCLPP_EP_COMBINE_TMA / - // MSCCLPP_EP_INTRA_ALLSENDER) and the input dtype. Shared by - // intranodeNotifyDispatch, intranodeDispatch and getIntranodeDispatchNumChannels - // so all three resolve identical channel counts. `xElementSize == 2` (BF16) - // is required for the all-sender path. - void computeIntranodeChannels(int xElementSize, const Config& config, int& dispatchNumSms, bool& allSender, - int& numChannels) const; - +class MoEHighThroughputRuntime { public: - MoEHighThroughputRuntime(int rank, int numRanks, int64_t numNvlBytes, int64_t numRdmaBytes); - + MoEHighThroughputRuntime(mscclpp::Communicator& communicator, int64_t maxHiddenBytes, const Config& config); ~MoEHighThroughputRuntime() noexcept(false); bool isAvailable() const; - bool isInternodeAvailable() const; - int getNumRdmaRanks() const; + void layout(int* numTokensPerRank, int* numTokensPerExpert, bool* isTokenInRank, const int64_t* topkIdx, + int numTokens, int numTopk, int numExperts, cudaStream_t stream); - int getRdmaRank() const; + int getDispatchNumChannels(int xElementSize) const; - int getRootRdmaRank(bool global) const; + void* resolveRecvXBuffer(int numTokens, int numRecvTokens, int hidden, int xElementSize) const; - int getLocalDeviceId() const; + int notifyDispatch(int* rankPrefixMatrix, int* channelPrefixMatrix, int* numRecvTokensPerExpert, + const int* numTokensPerRank, const int* numTokensPerExpert, const bool* isTokenInRank, + int numTokens, int numExperts, int xElementSize, int expertAlignment, cudaStream_t stream); - std::string getLocalIpcHandle() const; + void dispatch(void* recvX, float* recvXScales, int64_t* recvTopkIdx, float* recvTopkWeights, int* recvSrcIdx, + int* sendHead, int* recvChannelPrefixMatrix, const void* x, const float* xScales, + const int64_t* topkIdx, const float* topkWeights, const bool* isTokenInRank, + const int* rankPrefixMatrix, const int* channelPrefixMatrix, int numTokens, int hidden, int numTopk, + int numScales, int numExperts, int xElementSize, int numRecvTokens, bool cachedMode, + cudaStream_t stream); - std::string getLocalNvshmemUniqueId() const; + void combine(void* combinedX, float* combinedTopkWeights, const void* x, const float* topkWeights, const int* srcIdx, + const int* rankPrefixMatrix, const int* channelPrefixMatrix, const int* sendHead, int numTokens, + int numRecvTokens, int hidden, int numTopk, int xElementSize, int ringNumChannels, cudaStream_t stream); - mscclpp::UniqueId createUniqueId() const; + private: + void setup(mscclpp::Communicator& communicator); + void moveFifoSlots(int numSlots = 1); + void computeDispatchChannels(int xElementSize, int& dispatchNumSms, bool& allSender, int& numChannels) const; + bool canUseDirectRecvPool(int numTokens, int numRecvTokens, int hidden, int xElementSize) const; - void connect(mscclpp::UniqueId rootId); + int rank_; + int numRanks_; + int numNvlRanks_; + int numRanksPerIpcDomain_; + int head_ = 0; + int64_t maxHiddenBytes_; + size_t ringBufferBytes_ = 0; + size_t taskFifoOffset_ = 0; + size_t symmetricBufferBytes_ = 0; + size_t recvPoolBytes_ = 0; + bool available_ = false; + bool physicalRingBuffer_ = false; + bool directDispatchReady_ = false; + Config config_; - void sync(const std::vector& deviceIds, const std::vector>& allGatheredHandles, - const std::optional& rootUniqueIdOpt); + void* symmetricBuffer_ = nullptr; + void* recvPool_ = nullptr; + std::vector bufferPtrs_; + std::vector taskFifoPtrs_; + std::vector recvPoolPtrs_; + std::vector peerMemories_; + std::vector recvPoolMemories_; + void** bufferPtrsGpu_ = nullptr; + int** taskFifoPtrsGpu_ = nullptr; + void** recvPoolPtrsGpu_ = nullptr; + int* combineRecvIdxGpu_ = nullptr; - // Compute layout metadata for `topkIdx`. Outputs (caller pointers): - // numTokensPerRank [numRanks] (int) - // numTokensPerRdmaRank [numRdmaRanks] (int, may be nullptr when - // !isInternodeAvailable()) - // numTokensPerExpert [numExperts] (int) - // isTokenInRank [numTokens * numRanks] (bool) - void getDispatchLayout(int* numTokensPerRank, int* numTokensPerRdmaRank, int* numTokensPerExpert, bool* isTokenInRank, - const int64_t* topkIdx, int numTokens, int numTopk, int numExperts, cudaStream_t stream); - - // Number of channels the intranode dispatch resolves for this config + dtype. - // The caller uses it to size `channelPrefixMatrix` / `recvChannelPrefixMatrix` - // ([numRanks * numChannels]) and to pass `ringNumChannels` to `intranodeCombine`. - // Mirrors the env-driven channel-count logic in `intranodeDispatch` exactly. - int getIntranodeDispatchNumChannels(int xElementSize, const Config& config) const; - - // Resolve the recv-x output pointer for `intranodeDispatch`. When the - // zero-copy direct / all-sender path is active for this shape, returns the - // internal recv-pool view (recv_pool_local_ptr_ + header); the caller must - // pass this same pointer back as `recvX` (so expert outputs land in the pool - // for the TMA combine to gather). Returns nullptr when the caller should - // allocate `recvX` itself. Mirrors the `ep_intra_direct` gate in - // `intranodeDispatch` exactly so both agree. - void* resolveIntranodeRecvXBuffer(int numRecvTokens, int hidden, int xElementSize, const Config& config) const; - - // Resolve the recv-x output pointer for `internodeDispatch`. When the - // non-cached internode direct path is active for this shape, returns the - // internal recv-pool view (recv_pool_local_ptr_ + header); the caller must - // pass this same pointer back as `recvX` (so the cross-GPU forwarder writes - // hidden straight into the pool and the direct-gather combine reads it back). - // Returns nullptr when the caller should allocate `recvX` itself. Mirrors the - // `ep_use_direct` gate in `internodeDispatch` exactly (minus the cached-mode - // term, since the helper is only consulted on the non-cached forward path). - void* resolveInternodeRecvXBuffer(int numRecvTokens, int hidden, int xElementSize, const Config& config) const; - - // Number of channels the non-cached internode dispatch resolves for this - // config (mirrors `internodeDispatch`'s `ep_flat_dispatch_channels(num_sms)` - // under EP_DISPATCH_NCCLEP, else `num_sms/2`). The caller uses it to size the - // rdma/gbl channel-prefix matrices ([numRdmaRanks|numRanks * numChannels]). - int getInternodeDispatchNumChannels(const Config& config) const; - - // Per-token source-meta row width in bytes (recvSrcMeta is - // [numRecvTokens * getSourceMetaBytes()]). Exposes - // `internode::get_source_meta_bytes()` so the caller can size recvSrcMeta. - int getSourceMetaBytes() const; - - // Compile-time NUM_MAX_NVL_PEERS (4 on GB200, 8 on HGX). The caller uses it to - // size sendNvlHead ([numRdmaRecvTokens * NUM_MAX_NVL_PEERS]) and to validate - // combinedNvlHead's second dimension. - int getNumMaxNvlPeers() const; - - // Phase A (non-cached): send sizes / notify. Writes rankPrefixMatrix - // [numRanks*numRanks], channelPrefixMatrix [numRanks*numChannels] and the host - // array numRecvTokensPerExpert [numLocalExperts]; returns numRecvTokens. - // `xElementSize` (2 == BF16) is required to resolve the all-sender channel - // count identically to `intranodeDispatch`. - int intranodeNotifyDispatch(int* rankPrefixMatrix, int* channelPrefixMatrix, int* numRecvTokensPerExpert, - const int* numTokensPerRank, const int* numTokensPerExpert, const bool* isTokenInRank, - int numTokens, int numExperts, int xElementSize, int expertAlignment, - const Config& config, cudaStream_t stream); - - // Phase B (shared/cached tail): run the dispatch kernel. When `cachedMode` is - // true the caller supplies rankPrefixMatrix / channelPrefixMatrix / numRecvTokens - // from a prior dispatch and this method runs the cached barrier first; when - // false it assumes `intranodeNotifyDispatch` already ran. Writes the recv - // output pointers (recvXScales / recvTopkIdx / recvTopkWeights may be nullptr). - void intranodeDispatch(void* recvX, float* recvXScales, int64_t* recvTopkIdx, float* recvTopkWeights, int* recvSrcIdx, - int* sendHead, int* recvChannelPrefixMatrix, const void* x, const float* xScales, - const int64_t* topkIdx, const float* topkWeights, const bool* isTokenInRank, - const int* rankPrefixMatrix, const int* channelPrefixMatrix, int numTokens, int hidden, - int numTopk, int numScales, int numExperts, int xElementSize, int numRecvTokens, - bool cachedMode, const Config& config, cudaStream_t stream); - - // Combine (scatter-reduce). Output combinedX [numRecvTokens * hidden] where - // numRecvTokens == number of ORIGINAL tokens (== sendHead row count, known to - // the caller); combinedTopkWeights [numRecvTokens * numTopk] may be nullptr. - // `numTokens` is the dispatched token count (x row count); `ringNumChannels` - // is channelPrefixMatrix's column count (== getIntranodeDispatchNumChannels()). - void intranodeCombine(void* combinedX, float* combinedTopkWeights, const void* x, const float* topkWeights, - const int* srcIdx, const int* rankPrefixMatrix, const int* channelPrefixMatrix, - const int* sendHead, int numTokens, int numRecvTokens, int hidden, int numTopk, - int xElementSize, int ringNumChannels, const Config& config, cudaStream_t stream); - - // --------------------------------------------------------------------------- - // Internode (NVLink + RDMA) path. Declared torch-free; the bodies are - // throwing stubs in this drop and will be ported next (the intranode path is - // validated single-node first). See SIGNATURES.md for the notify->dispatch - // boundary contract. - // --------------------------------------------------------------------------- - - // Phase A (non-cached): writes rdmaChannelPrefixMatrix [numRdmaRanks*numChannels], - // recvRdmaRankPrefixSum [numRdmaRanks], gblChannelPrefixMatrix [numRanks*numChannels], - // recvGblRankPrefixSum [numRanks], host numRecvTokensPerExpert [numLocalExperts]; - // returns the primary count numRecvTokens and the secondary count via - // *numRdmaRecvTokens. - int internodeNotifyDispatch(int* rdmaChannelPrefixMatrix, int* recvRdmaRankPrefixSum, int* gblChannelPrefixMatrix, - int* recvGblRankPrefixSum, int* numRecvTokensPerExpert, int* numRdmaRecvTokens, - const int* numTokensPerRank, const int* numTokensPerRdmaRank, - const int* numTokensPerExpert, const bool* isTokenInRank, int numTokens, int numExperts, - int hidden, int numScales, int numTopk, int xElementSize, int expertAlignment, - const Config& config, cudaStream_t stream); - - // Phase B (shared/cached tail): runs the internode dispatch kernel. Writes the - // recv outputs (recvXScales / recvTopkIdx / recvTopkWeights / recvSrcMeta / - // recvRdmaChannelPrefixMatrix / recvGblChannelPrefixMatrix / sendRdmaHead / - // sendNvlHead may be nullptr in cached mode). - void internodeDispatch(void* recvX, float* recvXScales, int64_t* recvTopkIdx, float* recvTopkWeights, - void* recvSrcMeta, int* recvRdmaChannelPrefixMatrix, int* recvGblChannelPrefixMatrix, - int* sendRdmaHead, int* sendNvlHead, const void* x, const float* xScales, - const int64_t* topkIdx, const float* topkWeights, const bool* isTokenInRank, - const int* rdmaChannelPrefixMatrix, const int* recvRdmaRankPrefixSum, - const int* gblChannelPrefixMatrix, const int* recvGblRankPrefixSum, int numTokens, int hidden, - int numTopk, int numScales, int numExperts, int xElementSize, int numRecvTokens, - int numRdmaRecvTokens, bool cachedMode, const Config& config, cudaStream_t stream); - - // Combine. Output combinedX [numCombinedTokens * hidden] (numCombinedTokens == - // is_combined_token_in_rank row count, known to the caller); combinedTopkWeights - // [numCombinedTokens * numTopk] may be nullptr. - void internodeCombine(void* combinedX, float* combinedTopkWeights, const void* x, const float* topkWeights, - const void* srcMeta, const bool* isCombinedTokenInRank, const int* rdmaChannelPrefixMatrix, - const int* rdmaRankPrefixSum, const int* gblChannelPrefixMatrix, const int* combinedRdmaHead, - const int* combinedNvlHead, int numTokens, int numCombinedTokens, int hidden, int numTopk, - int xElementSize, const Config& config, cudaStream_t stream); + volatile int* moeRecvCounter_ = nullptr; + int* moeRecvCounterMapped_ = nullptr; + volatile int* moeRecvExpertCounter_ = nullptr; + int* moeRecvExpertCounterMapped_ = nullptr; }; } // namespace ep diff --git a/src/ext/ep/include/api.cuh b/src/ext/ep/include/api.cuh index 2f90e5eb..05390245 100644 --- a/src/ext/ep/include/api.cuh +++ b/src/ext/ep/include/api.cuh @@ -4,9 +4,7 @@ // Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), // branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. // -// Private host-callable API exposed by the EP CUDA kernels. One-to-one port of -// DeepEP `csrc/kernels/api.cuh` minus the NVSHMEM-only internode entrypoints, -// which are still to be migrated. +// Private host-callable API exposed by the EP CUDA kernels. #pragma once @@ -14,8 +12,6 @@ #include #include -#include -#include #include namespace mscclpp { @@ -38,12 +34,14 @@ enum class DispatchLayout { }; // =========================================================================== -// Archived HT intranode (NVLink) runtime barrier. -// Implementations live under `src/ext/ep/ht/` and are not compiled into the active -// `mscclpp_ep_cpp` target. +// High-throughput intranode kernels. // =========================================================================== namespace intranode { +void get_dispatch_layout(const int64_t* topk_idx, int* num_tokens_per_rank, int* num_tokens_per_expert, + bool* is_token_in_rank, int num_tokens, int num_topk, int num_ranks, int num_experts, + cudaStream_t stream); + void barrier(int** task_fifo_ptrs, int head, int rank, int num_ranks, cudaStream_t stream); void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, int num_ranks, @@ -98,90 +96,6 @@ void intranode_meta_drain(void* pool_base, int64_t meta_base, int num_recv_token } // namespace intranode -// =========================================================================== -// Archived internode (NVLink + RDMA) high-throughput kernels. Ported from DeepEP -// `csrc/kernels/internode.cu` on branch `chhwang/dev-atomic-add-cleanup`. The -// implementations live under `src/ext/ep/ht/` and are not compiled into the -// active `mscclpp_ep_cpp` target. -// =========================================================================== -namespace internode { - -int get_source_meta_bytes(); - -void get_dispatch_layout(const int64_t* topk_idx, int* num_tokens_per_rank, int* num_tokens_per_rdma_rank, - int* num_tokens_per_expert, bool* is_token_in_rank, int num_tokens, int num_topk, - int num_ranks, int num_experts, cudaStream_t stream); - -void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mapped, int num_ranks, - const int* num_tokens_per_rdma_rank, int* moe_recv_rdma_counter_mapped, - const int* num_tokens_per_expert, int* moe_recv_expert_counter_mapped, int num_experts, - const bool* is_token_in_rank, int num_tokens, int num_channels, int hidden_int4, int num_scales, - int num_topk, int expert_alignment, int* rdma_channel_prefix_matrix, - int* recv_rdma_rank_prefix_sum, int* gbl_channel_prefix_matrix, int* recv_gbl_rank_prefix_sum, - void* rdma_buffer_ptr, int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, - int num_max_nvl_chunked_recv_tokens, int** task_fifo_ptrs, int head, int rank, cudaStream_t stream, - int64_t num_rdma_bytes, int64_t num_nvl_bytes, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_mc_ptr, void* nvls_dev_ptr, - size_t nvls_off_barrier, size_t nvls_off_data, uint64_t nvls_epoch, int nvls_per_peer_bytes); - -void dispatch(void* recv_x, float* recv_x_scales, int64_t* recv_topk_idx, float* recv_topk_weights, void* recv_src_meta, - const void* x, const float* x_scales, const int64_t* topk_idx, const float* topk_weights, - int* send_rdma_head, int* send_nvl_head, int* recv_rdma_channel_prefix_matrix, - int* recv_gbl_channel_prefix_matrix, const int* rdma_channel_prefix_matrix, - const int* recv_rdma_rank_prefix_sum, const int* gbl_channel_prefix_matrix, - const int* recv_gbl_rank_prefix_sum, int num_tokens, int hidden_int4, int num_scales, int num_topk, - int num_experts, const bool* is_token_in_rank, void* rdma_buffer_ptr, - int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, - int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank, int num_ranks, - bool is_cached_dispatch, cudaStream_t stream, int num_channels, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_head_mc, void* nvls_head_dev, - void* nvls_tail_mc, void* nvls_tail_dev, void* const* peer_rdma_bases, - // Increment 4: per-peer base pointers of the VMM-allocated recv-output pool - // (non-null enables cross-GPU forwarder direct-write to recv_x; nullptr = legacy path). - void* const* recv_pool_ptrs = nullptr, - // Increment 5 (inc5): domain-wide recv-pool bases indexed by GLOBAL rank - // (sender direct-write under kEpDirect; nullptr = inactive). - void* const* recv_pool_global_ptrs = nullptr, - // Increment 5 combine-direct (Stage 1): per-(token, dst global rank) recv-pool - // slot index written by the sender; consumed by combine's gather path. - int* ep_combine_recv_idx = nullptr); - -// Increment 6 (kEpFlat): post-dispatch metadata drain. Copies per-token metadata -// the sender wrote into the destination pool's META region into the recv_* output -// tensors (topk rebased to this rank's local expert range). Launched on the comm -// stream right after `dispatch` when MSCCLPP_EP_FLAT is set. -void flat_meta_drain(void* pool_base, int64_t meta_base, int num_recv_tokens, void* recv_src_meta, float* recv_x_scales, - int64_t* recv_topk_idx, float* recv_topk_weights, int num_scales, int num_topk, int num_experts, - int num_ranks, int rank, int64_t meta_slot_bytes, cudaStream_t stream); - -void cached_notify(int hidden_int4, int num_scales, int num_topk_idx, int num_topk_weights, int num_ranks, - int num_channels, int num_combined_tokens, int* combined_rdma_head, - const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, int* combined_nvl_head, - void* rdma_buffer_ptr, int num_max_rdma_chunked_recv_tokens, void** buffer_ptrs, - int num_max_nvl_chunked_recv_tokens, int** task_fifo_ptrs, int head, int rank, cudaStream_t stream, - int64_t num_rdma_bytes, int64_t num_nvl_bytes, bool is_cached_dispatch, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_mc_ptr = nullptr, - void* nvls_dev_ptr = nullptr, size_t nvls_off_barrier = 0, uint64_t nvls_epoch = 0); - -void combine(cudaDataType_t type, void* combined_x, float* combined_topk_weights, const bool* is_combined_token_in_rank, - const void* x, const float* topk_weights, const int* combined_rdma_head, const int* combined_nvl_head, - const void* src_meta, const int* rdma_channel_prefix_matrix, const int* rdma_rank_prefix_sum, - const int* gbl_channel_prefix_matrix, int num_tokens, int num_combined_tokens, int hidden, int num_topk, - void* rdma_buffer_ptr, int num_max_rdma_chunked_send_tokens, int num_max_rdma_chunked_recv_tokens, - void** buffer_ptrs, int num_max_nvl_chunked_send_tokens, int num_max_nvl_chunked_recv_tokens, int rank, - int num_ranks, cudaStream_t stream, int num_channels, bool low_latency_mode, - mscclpp::PortChannelDeviceHandle* port_channel_handles, - mscclpp::MemoryChannelDeviceHandle* memory_channel_handles, void* nvls_head_mc, void* nvls_head_dev, - void* nvls_tail_mc, void* nvls_tail_dev, void* const* peer_rdma_bases, - // Increment 5 combine-direct: peer recv-pool bases + dispatch gather map - // (non-null + kEpDirect => combine gathers from pools; nullptr = legacy 2-hop). - void* const* recv_pool_global_ptrs = nullptr, const int* ep_combine_recv_idx = nullptr); - -} // namespace internode - // =========================================================================== // Low-latency kernels for RDMA and IPC paths. Ported from DeepEP // `csrc/kernels/internode_ll.cu` with NVSHMEM/IBGDA device ops replaced by diff --git a/src/ext/ep/include/constants.cuh b/src/ext/ep/include/constants.cuh index ef583aa4..1eeb51cc 100644 --- a/src/ext/ep/include/constants.cuh +++ b/src/ext/ep/include/constants.cuh @@ -8,15 +8,6 @@ #pragma once -// Maximum number of intra-node NVLink peers per RDMA rank. -// - 8 for H100 NVL8 / HGX-style nodes (DeepEP upstream default). -// - 4 for Azure GB200 NVL72 (4 GPUs per NUMA host). -// Configurable via the CMake cache var `MSCCLPP_EP_NUM_MAX_NVL_PEERS` -// (see `src/ext/ep/CMakeLists.txt`). Default keeps DeepEP-parity at 8. -#ifndef NUM_MAX_NVL_PEERS -#define NUM_MAX_NVL_PEERS 8 -#endif -#define NUM_MAX_RDMA_PEERS 20 #define NUM_MAX_FIFO_SLOTS 32768 #define NUM_WORKSPACE_BYTES (32 * 1024 * 1024) #define NUM_MAX_LOCAL_EXPERTS 1024 diff --git a/src/ext/ep/include/device_helpers.cuh b/src/ext/ep/include/device_helpers.cuh index ce117801..0885abeb 100644 --- a/src/ext/ep/include/device_helpers.cuh +++ b/src/ext/ep/include/device_helpers.cuh @@ -3,7 +3,6 @@ #pragma once #include -#include #include #include "exception.cuh" diff --git a/src/ext/ep/include/launch.cuh b/src/ext/ep/include/launch.cuh index 2ddfe1e7..cbae659f 100644 --- a/src/ext/ep/include/launch.cuh +++ b/src/ext/ep/include/launch.cuh @@ -32,28 +32,6 @@ } \ } while (false) -#define SWITCH_RDMA_RANKS(case_macro) \ - do { \ - switch (num_ranks / NUM_MAX_NVL_PEERS) { \ - case 2: \ - case_macro(2); \ - case 3: \ - case_macro(3); \ - case 4: \ - case_macro(4); \ - case 8: \ - case_macro(8); \ - case 16: \ - case_macro(16); \ - case 18: \ - case_macro(18); \ - case 20: \ - case_macro(20); \ - default: \ - EP_HOST_ASSERT(false and "Unsupported RDMA ranks"); \ - } \ - } while (false) - #define SWITCH_RANKS_WITH_DTYPE(dtype, case_macro) \ do { \ switch (num_ranks) { \ diff --git a/test/python/ep/CMakeLists.txt b/test/python/ep/CMakeLists.txt index 57ab3917..69d60a8c 100644 --- a/test/python/ep/CMakeLists.txt +++ b/test/python/ep/CMakeLists.txt @@ -7,7 +7,6 @@ # # Example: # cmake -S test/python/ep -B test/python/ep/build \ -# -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=8 \ # -DCMAKE_CUDA_ARCHITECTURES=90 # cmake --build test/python/ep/build -j 64 @@ -26,8 +25,6 @@ endif() get_filename_component(_default_mscclpp_src "${CMAKE_CURRENT_SOURCE_DIR}/../../.." ABSOLUTE) set(MSCCLPP_SRC "${_default_mscclpp_src}" CACHE PATH "mscclpp source tree") -set(MSCCLPP_EP_NUM_MAX_NVL_PEERS "8" CACHE STRING - "Compile-time NUM_MAX_NVL_PEERS for the EP kernels (8 for HGX, 4 for GB200)") find_package(MPI REQUIRED) find_package(CUDAToolkit REQUIRED) @@ -76,10 +73,7 @@ target_include_directories(mscclpp_ep_bench PRIVATE "${MSCCLPP_INSTALL_DIR}/include" "${CUPTI_INCLUDE_DIR}") -target_compile_definitions(mscclpp_ep_bench PRIVATE - MSCCLPP_USE_CUDA - EP_DISPATCH_NCCLEP - NUM_MAX_NVL_PEERS=${MSCCLPP_EP_NUM_MAX_NVL_PEERS}) +target_compile_definitions(mscclpp_ep_bench PRIVATE MSCCLPP_USE_CUDA) target_compile_options(mscclpp_ep_bench PRIVATE $<$:--expt-relaxed-constexpr> diff --git a/test/python/ep/test_internode_multirank.py b/test/python/ep/test_internode_multirank.py deleted file mode 100644 index 3569b392..00000000 --- a/test/python/ep/test_internode_multirank.py +++ /dev/null @@ -1,476 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -"""Multi-rank internode (HT) functional validation for mscclpp_ep. - -Launch on each node with (example: 2 nodes x 8 GPUs = 16 ranks): - - # on master (NODE_RANK=0): - MASTER_ADDR= MASTER_PORT=29600 NODE_RANK=0 \ - torchrun --nnodes=2 --nproc_per_node=8 \ - --rdzv-backend=c10d --rdzv-endpoint=:29600 \ - test/python/ep/test_internode_multirank.py - - # on worker (NODE_RANK=1): - MASTER_ADDR= MASTER_PORT=29600 NODE_RANK=1 \ - torchrun --nnodes=2 --nproc_per_node=8 \ - --rdzv-backend=c10d --rdzv-endpoint=:29600 \ - test/python/ep/test_internode_multirank.py - -Round-trip dispatch + combine using internode HT kernels across nodes. - -Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass -that times dispatch and combine **separately** with CUDA events. Reports -per-phase latency (max across ranks) plus aggregate effective bandwidth -(sum across ranks). Override iteration counts with -``MSCCLPP_EP_BENCH_WARMUP`` / ``MSCCLPP_EP_BENCH_ITERS`` and the bench -problem size with ``MSCCLPP_EP_BENCH_TOKENS`` / ``_HIDDEN``. -""" - -from __future__ import annotations - -import os -import sys - -# Disable ProcessGroupNCCL's HeartbeatMonitor before importing torch.distributed. -# It runs in a background thread polling the TCPStore; under mpirun, rank 0 -# (the store server) can exit before non-zero ranks finish teardown, producing -# noisy 'recvValue failed / Connection was likely closed' stack traces. -os.environ.setdefault("TORCH_NCCL_ENABLE_MONITORING", "0") - -import torch -import torch.distributed as dist - - -def _detect_local_world_size(): - """Number of GPUs per node (4 on GB200, 8 on H100/A100, etc.). - - Resolution order: - 1. `MSCCLPP_EP_LOCAL_WORLD_SIZE` env var (matches the C++ side). - 2. `LOCAL_WORLD_SIZE` (torchrun) or `OMPI_COMM_WORLD_LOCAL_SIZE` (mpirun). - 3. `torch.cuda.device_count()` on the current host. - """ - for var in ("MSCCLPP_EP_LOCAL_WORLD_SIZE", "LOCAL_WORLD_SIZE", "OMPI_COMM_WORLD_LOCAL_SIZE"): - v = os.environ.get(var) - if v and int(v) > 0: - return int(v) - return max(1, torch.cuda.device_count()) - - -def init_dist(): - rank = int(os.environ["RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - local_world_size = _detect_local_world_size() - local_rank = int(os.environ.get("LOCAL_RANK", rank % local_world_size)) - torch.cuda.set_device(local_rank) - dist.init_process_group( - backend="nccl", world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{local_rank}") - ) - return rank, world_size, local_rank, dist.new_group(list(range(world_size))) - - -def inplace_unique(x: torch.Tensor, num_slots: int): - assert x.dim() == 2 - mask = x < 0 - x_padded = x.masked_fill(mask, num_slots) - bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) - bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) - bin_count = bin_count[:, :num_slots] - sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) - sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) - sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values - x[:, :].fill_(-1) - valid_len = min(num_slots, x.size(1)) - x[:, :valid_len] = sorted_bin_idx[:, :valid_len] - - -def main(): - rank, num_ranks, local_rank, group = init_dist() - from mscclpp import CommGroup - import mscclpp.ep as ep - - ep_group = CommGroup(torch_group=group) - - NUM_MAX_NVL_PEERS = _detect_local_world_size() - assert ( - num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS - ), f"expected >1 node with {NUM_MAX_NVL_PEERS} GPUs each, got num_ranks={num_ranks}" - num_nodes = num_ranks // NUM_MAX_NVL_PEERS - num_local_ranks = NUM_MAX_NVL_PEERS - - # Small settings for functional check - import os as _os - - num_tokens = int(_os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "128")) - hidden = int(_os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "1024")) - num_topk = int(_os.environ.get("MSCCLPP_EP_BENCH_TOPK", str(min(4, num_ranks)))) - _experts_env = _os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", "") - num_experts = int(_experts_env) if _experts_env else num_ranks * 4 - assert num_experts % num_ranks == 0 - - torch.manual_seed(0xA1B2 + rank) - - scores = torch.randn((num_tokens, num_experts), device="cuda", dtype=torch.float32).abs() + 1 - topk_idx = torch.topk(scores, num_topk, dim=-1, sorted=False).indices - topk_weights = torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") - - rank_idx = topk_idx // (num_experts // num_ranks) - rank_idx.masked_fill_(topk_idx == -1, -1) - inplace_unique(rank_idx, num_ranks) - - rdma_rank_idx = rank_idx // num_local_ranks - rdma_rank_idx.masked_fill_(rank_idx == -1, -1) - inplace_unique(rdma_rank_idx, num_nodes) - - num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda") - for i in range(num_experts): - num_tokens_per_expert[i] = (topk_idx == i).sum() - - num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda") - num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda") - token_idx_in_rank = torch.full((num_ranks, num_tokens), -1, dtype=torch.long, device="cuda") - for i in range(num_ranks): - num_tokens_per_rank[i] = (rank_idx == i).sum() - token_sel = (rank_idx == i).max(dim=-1).values - cnt = token_sel.sum().item() - tokens = torch.sort(token_sel.to(torch.int), descending=True).indices - tokens[:cnt] = torch.sort(tokens[:cnt]).values - token_idx_in_rank[i][tokens[:cnt]] = torch.arange(cnt, dtype=torch.long, device="cuda") - for i in range(num_nodes): - num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum() - token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int) - is_token_in_rank = token_idx_in_rank >= 0 - - x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - - moe = ep.MoECommunicator( - comm=ep_group, - num_experts=num_experts, - hidden_size=hidden, - topk=num_topk, - max_tokens_per_rank=num_tokens, - mode=ep.MoEMode.HIGH_THROUGHPUT, - num_sms=int(os.environ.get("MSCCLPP_EP_NSM", "152")), - nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")), - nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")), - rdma_chunked_send=int(os.environ.get("MSCCLPP_EP_RDMA_SEND", "16")), - rdma_chunked_recv=int(os.environ.get("MSCCLPP_EP_RDMA_RECV", "128")), - ) - if rank == 0: - print( - f"[cfg] num_nodes={num_nodes} num_ranks={num_ranks} num_tokens={num_tokens} " - f"hidden={hidden} num_experts={num_experts} num_topk={num_topk}", - flush=True, - ) - - print( - f"[rank {rank}] MoECommunicator created is_available={moe.is_available()} " - f"is_internode={moe.is_internode_available()}", - flush=True, - ) - assert moe.is_available() and moe.is_internode_available() - assert moe.is_internode(), "expected the communicator to select the internode HT transport" - - dispatch_out, handle = moe.dispatch( - x, - topk_idx, - topk_weights, - ) - recv_x = dispatch_out.tokens - dist.barrier(group=group) - - assert recv_x.dim() == 2 and recv_x.size(1) == hidden - local_experts = num_experts // num_ranks - all_expert_counts = torch.empty((num_ranks, num_experts), dtype=num_tokens_per_expert.dtype, device="cuda") - dist.all_gather_into_tensor(all_expert_counts, num_tokens_per_expert, group=group) - expected_counts = all_expert_counts[:, rank * local_experts : (rank + 1) * local_experts].sum(dim=0).cpu().tolist() - assert dispatch_out.layout.num_tokens_per_expert is not None - actual_counts = [int(count) for count in dispatch_out.layout.num_tokens_per_expert] - assert actual_counts == [int(count) for count in expected_counts] - if rank == 0: - print(f"[dispatch] OK (recv {recv_x.size(0)} tokens)", flush=True) - - # Keep the existing dispatch/combine phase guard for internode HT until the - # backend wires a proper stream-dependency hand-off. - torch.cuda.synchronize() - dist.barrier(group=group) - - combined_x = moe.combine(recv_x, handle) - - num_dst = is_token_in_rank.sum(dim=1).to(torch.float32) - expected = num_dst * float(rank) - got = combined_x.float().mean(dim=1) - diff = (got - expected).abs().max().item() - max_exp = expected.abs().max().item() - print(f"[combine r{rank}] max|got-expected|={diff:.4e} max|expected|={max_exp:.4e}", flush=True) - # bf16 accumulator has 7-bit mantissa; intermediate partial sums can - # round at ulp = max_exp * 2**-7. Use a tolerance that scales with magnitude. - tol = max(1e-2, max_exp * (1.0 / 64)) - assert diff <= tol, f"rank{rank}: combine mismatch max diff {diff} > tol {tol} (max_exp={max_exp})" - - dist.barrier(group=group) - if rank == 0: - print("PASS", flush=True) - - # ------------------------------------------------------------------ - # Optional benchmark (enable with MSCCLPP_EP_BENCH=1). - # ------------------------------------------------------------------ - if os.environ.get("MSCCLPP_EP_BENCH", "0") != "1": - return - - warmup = int(os.environ.get("MSCCLPP_EP_BENCH_WARMUP", "5")) - iters = int(os.environ.get("MSCCLPP_EP_BENCH_ITERS", "20")) - bench_tokens = int(os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "4096")) - bench_hidden = int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168")) - # Allow overriding num_experts / num_topk for the bench phase to match - # NCCL-EP's `ep_bench -a ht` defaults (256 experts, top-8). The functional - # check above still uses the smaller (num_experts=num_ranks*4, topk=4) - # configuration. - bench_num_experts = int(os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", str(num_experts))) - bench_num_topk = int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", str(num_topk))) - if bench_num_experts % num_ranks != 0: - if rank == 0: - print( - f"[bench] skip: num_experts={bench_num_experts} not divisible " f"by num_ranks={num_ranks}", flush=True - ) - return - if bench_num_topk > bench_num_experts: - if rank == 0: - print(f"[bench] skip: topk={bench_num_topk} > experts={bench_num_experts}", flush=True) - return - - scores_b = torch.randn((bench_tokens, bench_num_experts), device="cuda", dtype=torch.float32).abs() + 1 - topk_idx_b = torch.topk(scores_b, bench_num_topk, dim=-1, sorted=False).indices - topk_weights_b = torch.ones((bench_tokens, bench_num_topk), dtype=torch.float32, device="cuda") - rank_idx_b = topk_idx_b // (bench_num_experts // num_ranks) - rank_idx_b.masked_fill_(topk_idx_b == -1, -1) - inplace_unique(rank_idx_b, num_ranks) - rdma_rank_idx_b = rank_idx_b // num_local_ranks - rdma_rank_idx_b.masked_fill_(rank_idx_b == -1, -1) - inplace_unique(rdma_rank_idx_b, num_nodes) - - num_tokens_per_expert_b = torch.zeros((bench_num_experts,), dtype=torch.int, device="cuda") - for i in range(bench_num_experts): - num_tokens_per_expert_b[i] = (topk_idx_b == i).sum() - num_tokens_per_rank_b = torch.empty((num_ranks,), dtype=torch.int, device="cuda") - num_tokens_per_rdma_rank_b = torch.empty((num_nodes,), dtype=torch.int, device="cuda") - token_idx_in_rank_b = torch.full((num_ranks, bench_tokens), -1, dtype=torch.long, device="cuda") - for i in range(num_ranks): - num_tokens_per_rank_b[i] = (rank_idx_b == i).sum() - token_sel = (rank_idx_b == i).max(dim=-1).values - cnt = token_sel.sum().item() - tokens = torch.sort(token_sel.to(torch.int), descending=True).indices - tokens[:cnt] = torch.sort(tokens[:cnt]).values - token_idx_in_rank_b[i][tokens[:cnt]] = torch.arange(cnt, dtype=torch.long, device="cuda") - for i in range(num_nodes): - num_tokens_per_rdma_rank_b[i] = (rdma_rank_idx_b == i).sum() - token_idx_in_rank_b = token_idx_in_rank_b.T.contiguous().to(torch.int) - is_token_in_rank_b = token_idx_in_rank_b >= 0 - x_b = torch.ones((bench_tokens, bench_hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - - # Drive the benchmark through the public high-level API. The communicator - # auto-selects internode HT when the RDMA size hint is non-zero. The first - # (uncached) dispatch records routing layout on the returned handle; - # subsequent dispatches reuse it via previous_handle, skipping host-side - # layout computation. This isolates the on-GPU dispatch-kernel cost - # (NCCL-EP ep_bench convention). - moe = ep.MoECommunicator( - comm=ep_group, - num_experts=bench_num_experts, - hidden_size=bench_hidden, - topk=bench_num_topk, - max_tokens_per_rank=bench_tokens, - mode=ep.MoEMode.HIGH_THROUGHPUT, - num_sms=int(os.environ.get("MSCCLPP_EP_NSM", "152")), - nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")), - nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")), - rdma_chunked_send=int(os.environ.get("MSCCLPP_EP_RDMA_SEND", "16")), - rdma_chunked_recv=int(os.environ.get("MSCCLPP_EP_RDMA_RECV", "128")), - ) - assert moe.is_available() and moe.is_internode_available() - assert moe.is_internode(), "expected the communicator to select the internode HT transport" - - # One uncached dispatch to build the cached routing layout on the handle. - _handle0 = moe.dispatch(x_b, topk_idx_b, topk_weights_b)[1] - - def _dispatch_cached(): - return moe.dispatch(x_b, topk_idx_b, topk_weights_b, previous_handle=_handle0) - - def _combine(dout): - dispatch_out_, handle_ = dout - moe.combine(dispatch_out_.tokens, handle_) - - # Warmup (full round-trip with the sync/barrier guard between phases, - # matching the correctness-path invariant: internode combine must observe - # the completed dispatch outputs before it launches). - for _ in range(warmup): - dout = _dispatch_cached() - torch.cuda.synchronize() - dist.barrier(group=group) - _combine(dout) - torch.cuda.synchronize() - dist.barrier(group=group) - - # Time dispatch alone (cached mode -- skips the host-side layout computation). - start_ev = torch.cuda.Event(enable_timing=True) - end_ev = torch.cuda.Event(enable_timing=True) - start_ev.record() - dout = None - for _ in range(iters): - dout = _dispatch_cached() - end_ev.record() - torch.cuda.synchronize() - disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters - - # Required guard before combine sees the dispatch outputs (see correctness - # path's XXX note). Not included in either phase's timing. - torch.cuda.synchronize() - dist.barrier(group=group) - - # Time combine alone (reusing the same dispatch output each iter). - start_ev.record() - for _ in range(iters): - _combine(dout) - end_ev.record() - torch.cuda.synchronize() - comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters - - # Per-rank "send bytes" matches NCCL-EP's `ep_bench` accounting (`RDMA_send`): - # bench_tokens * hidden * sizeof(bf16). Each rank ships its `bench_tokens` - # input rows out (some replicated to multiple peers); NCCL-EP normalizes by - # the input footprint, not by the recv-side fan-out. We use the same - # convention here so `per_rank_bw` is directly comparable across stacks. - bytes_one_way = bench_tokens * bench_hidden * x_b.element_size() - - # NCCL-EP `ep_bench` six-metric breakdown. - # Send-side accounting follows NCCL-EP: count unique (token, dst_node) pairs. - # `num_tokens_per_rdma_rank_b[n]` is exactly that count for node `n`. - # Recv-side accounting: each rank reports `num_tokens_per_rank_b[r]` - # (tokens it sends to dst rank `r`); an `all_to_all_single` lets every - # rank read how many tokens each source rank sent to it. - bytes_per_token = bench_hidden * x_b.element_size() - local_node = rank // num_local_ranks - nodes_unique = num_tokens_per_rdma_rank_b.to(torch.int64) - total_send_tokens_local = int(nodes_unique.sum().item()) - nvl_send_tokens_local = int(nodes_unique[local_node].item()) - rdma_send_tokens_local = total_send_tokens_local - nvl_send_tokens_local - # Replaced dist.all_to_all_single (NCCL socket transport fails with - # NCCL_IB_DISABLE=1 internode) with all_gather_into_tensor + transpose, - # which works on the same socket-NCCL setup the LL test uses. - _send_row = num_tokens_per_rank_b.to(torch.int64).contiguous() - _gathered = torch.empty(num_ranks * num_ranks, dtype=torch.int64, device="cuda") - dist.all_gather_into_tensor(_gathered, _send_row, group=group) - recv_from_src = _gathered.view(num_ranks, num_ranks)[:, rank].contiguous() - src_node = torch.arange(num_ranks, device="cuda") // num_local_ranks - remote_mask = (src_node != local_node).to(torch.int64) - total_recv_tokens_local = int(recv_from_src.sum().item()) - rdma_recv_tokens_local = int((recv_from_src * remote_mask).sum().item()) - - # Average per-rank token counts across ranks (matches NCCL-EP `Byte counts (per rank avg)`). - counts_t = torch.tensor( - [total_send_tokens_local, rdma_send_tokens_local, total_recv_tokens_local, rdma_recv_tokens_local], - dtype=torch.float64, - device="cuda", - ) - dist.all_reduce(counts_t, op=dist.ReduceOp.SUM, group=group) - counts_avg = (counts_t / num_ranks).tolist() - total_send_avg, rdma_send_avg, total_recv_avg, rdma_recv_avg = counts_avg - total_send_bytes = total_send_avg * bytes_per_token - rdma_send_bytes = rdma_send_avg * bytes_per_token - total_recv_bytes = total_recv_avg * bytes_per_token - rdma_recv_bytes = rdma_recv_avg * bytes_per_token - nvl_send_bytes = total_send_bytes - rdma_send_bytes - nvl_recv_bytes = total_recv_bytes - rdma_recv_bytes - - # Reduce timings: report min/avg/max and base BW on AVG to match NCCL-EP's - # `ep_bench.cu` convention. - disp_min_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda") - disp_avg_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda") - disp_max_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda") - comb_min_t = torch.tensor([comb_us], dtype=torch.float64, device="cuda") - comb_avg_t = torch.tensor([comb_us], dtype=torch.float64, device="cuda") - comb_max_t = torch.tensor([comb_us], dtype=torch.float64, device="cuda") - dist.all_reduce(disp_min_t, op=dist.ReduceOp.MIN, group=group) - dist.all_reduce(disp_avg_t, op=dist.ReduceOp.SUM, group=group) - dist.all_reduce(disp_max_t, op=dist.ReduceOp.MAX, group=group) - dist.all_reduce(comb_min_t, op=dist.ReduceOp.MIN, group=group) - dist.all_reduce(comb_avg_t, op=dist.ReduceOp.SUM, group=group) - dist.all_reduce(comb_max_t, op=dist.ReduceOp.MAX, group=group) - disp_avg_us = disp_avg_t.item() / num_ranks - comb_avg_us = comb_avg_t.item() / num_ranks - disp_bw_per_rank = bytes_one_way / (disp_avg_us * 1e-6) / 1e9 - comb_bw_per_rank = bytes_one_way / (comb_avg_us * 1e-6) / 1e9 - # Six-metric BW (NCCL-EP convention). Combine reverses send<->recv: - # in combine, this rank pushes back what it received in dispatch. - disp_t_s = disp_avg_us * 1e-6 - comb_t_s = comb_avg_us * 1e-6 - d_send_total_bw = total_send_bytes / disp_t_s / 1e9 - d_send_nvl_bw = nvl_send_bytes / disp_t_s / 1e9 - d_send_rdma_bw = rdma_send_bytes / disp_t_s / 1e9 - d_recv_total_bw = total_recv_bytes / disp_t_s / 1e9 - d_recv_nvl_bw = nvl_recv_bytes / disp_t_s / 1e9 - d_recv_rdma_bw = rdma_recv_bytes / disp_t_s / 1e9 - c_send_total_bw = total_recv_bytes / comb_t_s / 1e9 - c_send_nvl_bw = nvl_recv_bytes / comb_t_s / 1e9 - c_send_rdma_bw = rdma_recv_bytes / comb_t_s / 1e9 - c_recv_total_bw = total_send_bytes / comb_t_s / 1e9 - c_recv_nvl_bw = nvl_send_bytes / comb_t_s / 1e9 - c_recv_rdma_bw = rdma_send_bytes / comb_t_s / 1e9 - if rank == 0: - print( - f"[bench internode HT] nodes={num_nodes} num_ranks={num_ranks} " - f"tokens={bench_tokens} hidden={bench_hidden} " - f"experts={bench_num_experts} topk={bench_num_topk} " - f"warmup={warmup} iters={iters}", - flush=True, - ) - print( - f" dispatch: avg={disp_avg_us:.1f}us min={disp_min_t.item():.1f}us max={disp_max_t.item():.1f}us " - f"per_rank_bw={disp_bw_per_rank:.2f} GB/s " - f"agg_bw={disp_bw_per_rank * num_ranks:.2f} GB/s (BW @ avg time)", - flush=True, - ) - print( - f" send: total={d_send_total_bw:.2f} nvl={d_send_nvl_bw:.2f} rdma={d_send_rdma_bw:.2f} GB/s " - f"recv: total={d_recv_total_bw:.2f} nvl={d_recv_nvl_bw:.2f} rdma={d_recv_rdma_bw:.2f} GB/s", - flush=True, - ) - print( - f" combine : avg={comb_avg_us:.1f}us min={comb_min_t.item():.1f}us max={comb_max_t.item():.1f}us " - f"per_rank_bw={comb_bw_per_rank:.2f} GB/s " - f"agg_bw={comb_bw_per_rank * num_ranks:.2f} GB/s (BW @ avg time)", - flush=True, - ) - print( - f" send: total={c_send_total_bw:.2f} nvl={c_send_nvl_bw:.2f} rdma={c_send_rdma_bw:.2f} GB/s " - f"recv: total={c_recv_total_bw:.2f} nvl={c_recv_nvl_bw:.2f} rdma={c_recv_rdma_bw:.2f} GB/s", - flush=True, - ) - print( - f" byte counts (per rank avg): " - f"total_send={total_send_bytes/1e6:.2f} MB ({total_send_avg:.0f} tok) " - f"rdma_send={rdma_send_bytes/1e6:.2f} MB ({rdma_send_avg:.0f} tok) " - f"total_recv={total_recv_bytes/1e6:.2f} MB ({total_recv_avg:.0f} tok) " - f"rdma_recv={rdma_recv_bytes/1e6:.2f} MB ({rdma_recv_avg:.0f} tok)", - flush=True, - ) - - -if __name__ == "__main__": - try: - main() - except Exception: - import traceback - - traceback.print_exc() - sys.exit(1) - finally: - # Ordered shutdown: barrier so every rank reaches teardown before the - # TCPStore server (rank 0) exits, then destroy the PG. Without this, - # ProcessGroupNCCL's HeartbeatMonitor on non-zero ranks logs noisy - # "recvValue failed / Connection was likely closed" stack traces. - if dist.is_initialized(): - try: - dist.barrier() - except Exception: - pass - dist.destroy_process_group() diff --git a/test/python/ep/test_intranode_multirank.py b/test/python/ep/test_intranode_multirank.py index 977284df..54869aeb 100644 --- a/test/python/ep/test_intranode_multirank.py +++ b/test/python/ep/test_intranode_multirank.py @@ -1,13 +1,13 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Multi-rank intranode functional validation for mscclpp_ep. +"""Multi-rank direct-fabric HT functional validation for mscclpp_ep. Launch with: torchrun --nproc_per_node= test/python/ep/test_intranode_multirank.py -Tests that the high-level ``MoECommunicator`` succeeds across N GPUs on a single -node and that a round-trip dispatch + combine preserves data (sum of top-k -weighted copies). +Tests that the high-level ``MoECommunicator`` succeeds across GPUs in one +detected GPU IPC/NVL fabric domain, including domains that span hosts, and that +a round-trip dispatch + combine preserves data. Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass that times dispatch and combine **separately** with CUDA events and @@ -127,6 +127,10 @@ def main(): ) print(f"[rank {rank}] MoECommunicator created is_available={moe.is_available()}", flush=True) assert moe.is_available() + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", str(num_ranks))) + expected_internode = num_ranks > local_world_size + assert moe.is_internode_available() == expected_internode + assert moe.is_internode() == expected_internode dispatch_out, handle = moe.dispatch( x, @@ -271,16 +275,13 @@ def main(): torch.cuda.synchronize() comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters - # Per-rank "send bytes" matches NCCL-EP's `ep_bench` accounting (`RDMA_send`): + # Per-rank "send bytes" matches NCCL-EP's `ep_bench` accounting: # bench_tokens * hidden * sizeof(bf16). Each rank ships its `bench_tokens` # input rows out (some replicated to multiple peers); NCCL-EP normalizes by # the input footprint, not by the recv-side fan-out. We use the same # convention here so `per_rank_bw` is directly comparable across stacks. bytes_one_way = bench_tokens * bench_hidden * x_b.element_size() - # NCCL-EP `ep_bench` six-metric breakdown - # (intranode -> single node, so rdma_*=0; nvl_*=total_*). - # # Send side follows NCCL-EP: count unique (token, dst_node) pairs. With a # single node every selected destination collapses to that node, so a # token with at least one valid expert contributes exactly one to @@ -288,32 +289,23 @@ def main(): # landing on this rank. bytes_per_token = bench_hidden * x_b.element_size() total_send_tokens_local = int(is_token_in_rank_b.any(dim=1).sum().item()) - rdma_send_tokens_local = 0 # intranode: no remote nodes - # Replaced dist.all_to_all_single (NCCL socket transport fails with - # NCCL_IB_DISABLE=1 internode) with all_gather_into_tensor + transpose, - # which works on the same socket-NCCL setup the LL test uses. _send_row = num_tokens_per_rank_b.to(torch.int64).contiguous() _gathered = torch.empty(num_ranks * num_ranks, dtype=torch.int64, device="cuda") dist.all_gather_into_tensor(_gathered, _send_row, group=group) recv_from_src = _gathered.view(num_ranks, num_ranks)[:, rank].contiguous() total_recv_tokens_local = int(recv_from_src.sum().item()) - rdma_recv_tokens_local = 0 # intranode # Average per-rank token counts across ranks (matches NCCL-EP `Byte counts (per rank avg)`). counts_t = torch.tensor( - [total_send_tokens_local, rdma_send_tokens_local, total_recv_tokens_local, rdma_recv_tokens_local], + [total_send_tokens_local, total_recv_tokens_local], dtype=torch.float64, device="cuda", ) dist.all_reduce(counts_t, op=dist.ReduceOp.SUM, group=group) counts_avg = (counts_t / num_ranks).tolist() - total_send_avg, rdma_send_avg, total_recv_avg, rdma_recv_avg = counts_avg + total_send_avg, total_recv_avg = counts_avg total_send_bytes = total_send_avg * bytes_per_token - rdma_send_bytes = rdma_send_avg * bytes_per_token total_recv_bytes = total_recv_avg * bytes_per_token - rdma_recv_bytes = rdma_recv_avg * bytes_per_token - nvl_send_bytes = total_send_bytes - rdma_send_bytes - nvl_recv_bytes = total_recv_bytes - rdma_recv_bytes # Reduce timings: report min/avg/max and base BW on AVG to match NCCL-EP's # `ep_bench.cu` convention. @@ -333,21 +325,12 @@ def main(): comb_avg_us = comb_avg_t.item() / num_ranks disp_bw_per_rank = bytes_one_way / (disp_avg_us * 1e-6) / 1e9 comb_bw_per_rank = bytes_one_way / (comb_avg_us * 1e-6) / 1e9 - # Six-metric BW (NCCL-EP convention). Combine reverses send<->recv. disp_t_s = disp_avg_us * 1e-6 comb_t_s = comb_avg_us * 1e-6 d_send_total_bw = total_send_bytes / disp_t_s / 1e9 - d_send_nvl_bw = nvl_send_bytes / disp_t_s / 1e9 - d_send_rdma_bw = rdma_send_bytes / disp_t_s / 1e9 d_recv_total_bw = total_recv_bytes / disp_t_s / 1e9 - d_recv_nvl_bw = nvl_recv_bytes / disp_t_s / 1e9 - d_recv_rdma_bw = rdma_recv_bytes / disp_t_s / 1e9 c_send_total_bw = total_recv_bytes / comb_t_s / 1e9 # combine sends back what dispatch received - c_send_nvl_bw = nvl_recv_bytes / comb_t_s / 1e9 - c_send_rdma_bw = rdma_recv_bytes / comb_t_s / 1e9 c_recv_total_bw = total_send_bytes / comb_t_s / 1e9 # combine receives back what dispatch sent - c_recv_nvl_bw = nvl_send_bytes / comb_t_s / 1e9 - c_recv_rdma_bw = rdma_send_bytes / comb_t_s / 1e9 if rank == 0: print( f"[bench intranode HT] tokens={bench_tokens} hidden={bench_hidden} " @@ -362,8 +345,7 @@ def main(): flush=True, ) print( - f" send: total={d_send_total_bw:.2f} nvl={d_send_nvl_bw:.2f} rdma={d_send_rdma_bw:.2f} GB/s " - f"recv: total={d_recv_total_bw:.2f} nvl={d_recv_nvl_bw:.2f} rdma={d_recv_rdma_bw:.2f} GB/s", + f" send={d_send_total_bw:.2f} GB/s recv={d_recv_total_bw:.2f} GB/s", flush=True, ) print( @@ -373,16 +355,13 @@ def main(): flush=True, ) print( - f" send: total={c_send_total_bw:.2f} nvl={c_send_nvl_bw:.2f} rdma={c_send_rdma_bw:.2f} GB/s " - f"recv: total={c_recv_total_bw:.2f} nvl={c_recv_nvl_bw:.2f} rdma={c_recv_rdma_bw:.2f} GB/s", + f" send={c_send_total_bw:.2f} GB/s recv={c_recv_total_bw:.2f} GB/s", flush=True, ) print( f" byte counts (per rank avg): " f"total_send={total_send_bytes/1e6:.2f} MB ({total_send_avg:.0f} tok) " - f"rdma_send={rdma_send_bytes/1e6:.2f} MB ({rdma_send_avg:.0f} tok) " - f"total_recv={total_recv_bytes/1e6:.2f} MB ({total_recv_avg:.0f} tok) " - f"rdma_recv={rdma_recv_bytes/1e6:.2f} MB ({rdma_recv_avg:.0f} tok)", + f"total_recv={total_recv_bytes/1e6:.2f} MB ({total_recv_avg:.0f} tok)", flush=True, )