From 57ea3dd5c9ffb26c87ef292518249d2f001666ab Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Mon, 29 Jun 2026 13:32:51 -0700 Subject: [PATCH] MoE Commnucator design doc (#818) Add API doc for MoE communication Refactor EP API for Low latency mode --- docs/quickstart.md | 10 +- include/mscclpp/gpu_data_types.hpp | 6 + pyproject.toml | 1 + python/mscclpp/_core/comm.py | 15 +- python/mscclpp/ext/ep/README.md | 756 ++++++++++++++++ python/mscclpp/ext/ep/__init__.py | 28 +- python/mscclpp/ext/ep/buffer.py | 193 ---- python/mscclpp/ext/ep/communicator.py | 536 +++++++++++ python/test/test_gpu_buffer_pool_nvls_zero.py | 135 +++ src/ext/ep/CMakeLists.txt | 94 +- src/ext/ep/README.md | 137 ++- src/ext/ep/bindings.cpp | 139 +-- src/ext/ep/config.hpp | 215 ++--- src/ext/ep/{ => ht}/buffer.cc | 273 +++--- src/ext/ep/{ => ht}/buffer.hpp | 30 +- src/ext/ep/ht/config.hpp | 106 +++ src/ext/ep/{ => ht}/event.hpp | 0 src/ext/ep/{ => ht}/kernels/buffer.cuh | 4 +- src/ext/ep/{ => ht}/kernels/internode.cu | 71 +- .../ep/{ => ht}/kernels/internode_ncclep.cuh | 62 +- .../ep/{ => ht}/kernels/intranode_kernel.cu | 15 +- src/ext/ep/{ => ht}/kernels/runtime.cu | 8 +- src/ext/ep/kernels/api.cuh | 197 +++- src/ext/ep/kernels/common.cuh | 96 ++ src/ext/ep/kernels/configs.cuh | 7 +- src/ext/ep/kernels/internode_ll.cu | 749 ---------------- src/ext/ep/kernels/launch.cuh | 16 - src/ext/ep/kernels/low_latency.cu | 842 ++++++++++++++++++ src/ext/ep/kernels/quantization.cuh | 112 +++ src/ext/ep/kernels/utils.cuh | 20 +- src/ext/ep/moe_runtime.cc | 367 ++++++++ src/ext/ep/moe_runtime.hpp | 82 ++ .../python/ext/ep/test_internode_multirank.py | 31 +- .../python/ext/ep/test_intranode_multirank.py | 29 +- .../ext/ep/test_low_latency_multirank.py | 199 ++--- 35 files changed, 3863 insertions(+), 1718 deletions(-) create mode 100644 python/mscclpp/ext/ep/README.md delete mode 100644 python/mscclpp/ext/ep/buffer.py create mode 100644 python/mscclpp/ext/ep/communicator.py create mode 100644 python/test/test_gpu_buffer_pool_nvls_zero.py rename src/ext/ep/{ => ht}/buffer.cc (90%) rename src/ext/ep/{ => ht}/buffer.hpp (93%) create mode 100644 src/ext/ep/ht/config.hpp rename src/ext/ep/{ => ht}/event.hpp (100%) rename src/ext/ep/{ => ht}/kernels/buffer.cuh (98%) rename src/ext/ep/{ => ht}/kernels/internode.cu (98%) rename src/ext/ep/{ => ht}/kernels/internode_ncclep.cuh (95%) rename src/ext/ep/{ => ht}/kernels/intranode_kernel.cu (99%) rename src/ext/ep/{ => ht}/kernels/runtime.cu (88%) create mode 100644 src/ext/ep/kernels/common.cuh delete mode 100644 src/ext/ep/kernels/internode_ll.cu create mode 100644 src/ext/ep/kernels/low_latency.cu create mode 100644 src/ext/ep/kernels/quantization.cuh create mode 100644 src/ext/ep/moe_runtime.cc create mode 100644 src/ext/ep/moe_runtime.hpp diff --git a/docs/quickstart.md b/docs/quickstart.md index 716fcf61..6230723f 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -105,6 +105,8 @@ Python 3.10 or later is required. ```bash # For NVIDIA platforms (specify your CUDA version) $ python -m pip install ".[cuda12]" +# For NVIDIA platforms with the Expert Parallel (MoE dispatch/combine) extension +$ python -m pip install ".[cuda12,ep]" # For AMD platforms $ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm6]" ``` @@ -116,14 +118,19 @@ $ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm6]" Optional extras can be installed by specifying them in brackets. Available extras: - **`cuda11`**, **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package for your CUDA version. - **`rocm6`**: Install CuPy from source for AMD ROCm platforms. +- **`ep`**: Build and install the Expert Parallel extension (`mscclpp.ext.ep`). The extension itself does + not add a PyTorch dependency, but the high-level Python API expects user-provided `torch.Tensor` inputs. + CUDA architectures 90 or newer are required for EP kernels. - **`benchmark`**: Install benchmark dependencies (mpi4py, prettytable, netifaces, matplotlib). - **`test`**: Install test dependencies (pytest, mpi4py, netifaces). ```bash # Example: install with CUDA 12 and benchmark extras $ python -m pip install ".[cuda12,benchmark]" +# Example: install with CUDA 12 and the EP extension +$ python -m pip install ".[cuda12,ep]" # Example: install with all extras for testing on CUDA 12 -$ python -m pip install ".[cuda12,benchmark,test]" +$ python -m pip install ".[cuda12,ep,benchmark,test]" ``` (mrc-support)= @@ -291,4 +298,3 @@ Version: 0.8.0.post1.dev0+gc632fee37.d20251007 mscclpp.version {'version': '0.8.0.post1.dev0+gc632fee37.d20251007', 'git_commit': 'g50382c567'} ``` - diff --git a/include/mscclpp/gpu_data_types.hpp b/include/mscclpp/gpu_data_types.hpp index 672434f9..b9d7191a 100644 --- a/include/mscclpp/gpu_data_types.hpp +++ b/include/mscclpp/gpu_data_types.hpp @@ -64,7 +64,13 @@ using __fp8x4_e5m2 = __nv_fp8x4_e5m2; #define __FP8_TYPES_EXIST__ #endif +// GCC's x86 intrinsic headers may define __bfloat16 before this header is +// included by host-only Torch extension translation units. Avoid redefining it +// there; CUDA device compilation still uses __nv_bfloat16. +#if !defined(MSCCLPP_AVOID_BFLOAT16_ALIAS) && !defined(__AVX512BF16INTRIN_H_INCLUDED) && \ + !defined(__AVX512BF16VLINTRIN_H_INCLUDED) using __bfloat16 = __nv_bfloat16; +#endif using __bfloat162 = __nv_bfloat162; #endif diff --git a/pyproject.toml b/pyproject.toml index 816052b5..6834ab7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ cuda11 = ["cupy-cuda11x"] cuda12 = ["cupy-cuda12x"] cuda13 = ["cupy-cuda13x"] rocm6 = ["cupy"] +ep = [] benchmark = [ "mpi4py", "prettytable", diff --git a/python/mscclpp/_core/comm.py b/python/mscclpp/_core/comm.py index d42349dd..8e05ca16 100644 --- a/python/mscclpp/_core/comm.py +++ b/python/mscclpp/_core/comm.py @@ -54,14 +54,21 @@ class CommGroup: import torch import torch.distributed as dist + backend = str(dist.get_backend(torch_group)).lower() + device = torch.device("cuda", torch.cuda.current_device()) if "nccl" in backend else torch.device("cpu") if rank == 0: - uniq_id_global = uniq_id pickled_data = pickle.dumps(uniq_id) - data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone() + size_tensor = torch.tensor([len(pickled_data)], dtype=torch.int64, device=device) else: - data_tensor = torch.zeros(256, dtype=torch.uint8) + size_tensor = torch.zeros(1, dtype=torch.int64, device=device) + dist.broadcast(size_tensor, src=0, group=torch_group) + payload_size = int(size_tensor.item()) + if rank == 0: + data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone().to(device) + else: + data_tensor = torch.zeros(payload_size, dtype=torch.uint8, device=device) dist.broadcast(data_tensor, src=0, group=torch_group) - uniq_id_global = pickle.loads(data_tensor.numpy().tobytes()) + uniq_id_global = pickle.loads(data_tensor.cpu().numpy().tobytes()) self.bootstrap.initialize(uniq_id_global) elif not interfaceIpPortTrio == "": assert rank >= 0 and size >= 1 diff --git a/python/mscclpp/ext/ep/README.md b/python/mscclpp/ext/ep/README.md new file mode 100644 index 00000000..40c066c2 --- /dev/null +++ b/python/mscclpp/ext/ep/README.md @@ -0,0 +1,756 @@ +# Expert Parallel Python API Design + +This document proposes a simplified Python API for MoE expert-parallel +dispatch and combine. It is a design draft for review, not a committed API +contract. + +## Goals + +The API should expose the tensors that the model naturally owns: + +- token activations, +- top-k expert ids, +- routing weights, +- quantization scales. + +It should adapt to different MoE runner backends, such as Triton ragged/grouped +GEMM, CUTLASS-style grouped GEMM, DeepGEMM, FlashInfer/CuteDSL, and custom MLP +kernels, without forcing one physical layout on every backend. + +The dispatch output should make the local MLP contract explicit: + +- the token layout returned by dispatch, +- per-local-expert valid token counts, +- optional expert offsets for flat layouts, +- optional quantization scale layout, +- optional overlap capability when the selected runner can notify combine. + +## Public class name + +Use `MoECommunicator` as the public class name: + +```python +from mscclpp.ext.ep import MoECommunicator + +moe_comm = MoECommunicator(...) +``` + +The class owns MoE dispatch/combine communication, but it does not own the MLP +compute backend. + +## MoECommunicator configuration + +`MoECommunicator` owns communication setup, scratch buffers, expert placement, +layout choice, and optional overlap resources. These fields should be configured +once instead of being passed to every dispatch/combine call. + +```python +@dataclass +class MoECommunicatorConfig: + # Communication + comm: Optional[mscclpp.CommGroup] = None + device: Optional[torch.device | int] = None + + # Expert topology + num_experts: int = 0 + num_local_experts: Optional[int] = None # inferred for even contiguous placement + local_expert_start: Optional[int] = None # inferred as rank * num_local_experts + + # Model shape and capacity + hidden_size: int = 0 + topk: int = 0 + max_tokens_per_rank: int = 0 + max_recv_tokens_per_rank: Optional[int] = None + + # Runtime mode and output layout + mode: MoEMode = MoEMode.LOW_LATENCY + output_layout: Optional[DispatchLayout] = None # default is derived from mode + + # Quantization defaults + input_dtype: Optional[torch.dtype] = None + quant_format: Optional[str] = None + + # Transport resources + num_rdma_qps_per_rank: int = 12 # RDMA QPs per peer rank; advanced tuning + num_sms: int = 20 + + # Overlap + enable_overlap: bool = False +``` + +The constructor can accept either a config object or keyword arguments: + +```python +moe_comm = MoECommunicator( + comm=comm, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + topk=topk, + max_tokens_per_rank=max_tokens, + mode=MoEMode.HIGH_THROUGHPUT, +) +``` + +### Communication fields + +`comm` is the MSCCL++ communication object used for rank information, +out-of-band metadata exchange, and connection setup. + +The class should cache: + +| Field | Purpose | +|---|---| +| `comm` | MSCCL++ communicator or `CommGroup` | +| `rank`, `world_size` | global EP rank information | +| `local_rank`, `device` | CUDA device binding | +| internal runtime | nanobind/C++ EP runtime implementation | + +### Expert placement fields + +The class needs enough information to map global expert ids to local expert +ids: + +```python +local_expert_id = global_expert_id - local_expert_start +``` + +For the common contiguous placement: + +```python +if num_local_experts is None: + assert num_experts % world_size == 0 + num_local_experts = num_experts // world_size + +if local_expert_start is None: + local_expert_start = rank * num_local_experts +``` + +So both fields may be `None` for the common even contiguous placement. If the +placement is uneven or non-contiguous, the caller must provide enough placement +information explicitly. The first version can require contiguous local experts; +a later version can add an explicit `expert_map` for arbitrary placement. + +### Runtime fields + +`MoECommunicator` should keep these runtime decisions internally: + +| Field | Purpose | +|---|---| +| `mode` | Backend selection (`"ll"` active; `"ht"` archived/not compiled) | +| `output_layout` | MLP input layout returned by dispatch | +| `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 | +| `dispatch_config`, `combine_config` | backend-specific tuning configs | +| `overlap_capability` | whether selected MLP/backend supports notify | + +The user should not pass these to `dispatch` unless explicitly overriding a +specialized advanced path. + +### Mode selection + +The active implementation supports `mode=MoEMode.LOW_LATENCY`. `mode` must be a +`MoEMode` enum value, not a string. `MoEMode.HIGH_THROUGHPUT` raises +`NotImplementedError` because the HT implementation is archived under +`src/ext/ep/ht/` and is not compiled into `mscclpp_ep_cpp`. + +```python +moe_comm = MoECommunicator(..., mode=MoEMode.LOW_LATENCY) +``` + +This keeps `MoECommunicator` policy-free. Serving frameworks such as SGLang can +choose a mode based on their own scheduling policy, batch shape, runner backend, +and benchmarking data once multiple active backends are available. + +The selected mode determines the default dispatch output layout: + +| Mode | Default layout | +|---|---| +| `ht` | `DispatchLayout.FLAT` | +| `ll` | `DispatchLayout.EXPERT_MAJOR` | + +`output_layout` may still be kept as an advanced override if a backend supports +multiple layouts within the same mode. + +Use `DispatchLayout` instead of string literals for this field: + +| Layout enum | Tensor shape | +|---|---| +| `DispatchLayout.FLAT` | HT: `[total_recv_tokens, hidden]`; LL: `[num_local_experts * max_slots_per_expert, hidden]` | +| `DispatchLayout.EXPERT_MAJOR` | `[num_local_experts, max_slots_per_expert, hidden]` | + +## MoECommunicator methods + +```python +class MoECommunicator: + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor] = None, + scales: Optional[QuantScales] = None, + *, + output_buffer: torch.Tensor, + stream: Optional[torch.cuda.Stream] = None, + ) -> tuple[DispatchOutput, DispatchHandle]: + ... + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor] = None, + stream: Optional[torch.cuda.Stream] = None, + ) -> torch.Tensor: + ... + + def dispatch_async(..., overlap_config: Optional[CommOverlapConfig] = None) -> DispatchRequest: + ... + + def combine_async(..., overlap_config: Optional[CommOverlapConfig] = None) -> CombineRequest: + ... + + def create_overlap_config( + self, + op: str, # "dispatch" or "combine" + *, + handle: Optional[DispatchHandle] = None, + level: str = "op", # "op" or "block" + ) -> CommOverlapConfig: + ... +``` + +The blocking `dispatch` and `combine` methods should be the default path. The +`*_async` methods and `create_overlap_config` are optional advanced APIs for +communication/computation overlap. +If `stream` is not provided, both methods launch on `torch.cuda.current_stream()`. + +## High-level API + +```python +dispatch_out, handle = moe_comm.dispatch( + input, + topk_ids, + weights=None, + scales=None, + output_buffer=output_buffer, +) + +expert_output = mlp( + dispatch_out.tokens, + dispatch_out.num_tokens_per_expert, + dispatch_out.scales, +) + +output = moe_comm.combine(expert_output, handle) +``` + +`dispatch_out` is for the local MLP. `handle` is for `combine`. The MLP should +not need to inspect the opaque handle. + +## Proposed types + +```python +@dataclass +class QuantScales: + local: Optional[torch.Tensor] = None + global_scale: Optional[torch.Tensor] = None + format: Optional[str] = None + block_size: Optional[int] = None + + +class DispatchLayout(str, Enum): + FLAT = "flat" + EXPERT_MAJOR = "expert_major" + + +@dataclass +class DispatchOutput: + tokens: torch.Tensor + scales: Optional[QuantScales] + num_tokens_per_expert: torch.Tensor | list[int] + expert_offsets: Optional[torch.Tensor] = None + layout: DispatchLayout = DispatchLayout.FLAT + + +class DispatchHandle: + """Opaque handle returned by dispatch and consumed by combine.""" + + +@dataclass +class CommOverlapConfig: + op: str # "dispatch" or "combine" + level: str = "op" # "op" or "block" + stream: Optional[torch.cuda.Stream] = None + wait_event: Optional[torch.cuda.Event] = None + signal: Optional[torch.Tensor] = None + num_comm_sms: Optional[int] = None + block_m: Optional[int] = None + block_ready_value: Optional[int] = None + +``` + +`create_overlap_config` creates optional overlap configuration for async +dispatch/combine calls. + +```python +dispatch_overlap_config = moe_comm.create_overlap_config(op="dispatch") +combine_overlap_config = moe_comm.create_overlap_config(op="combine", handle=handle) +``` + +Operation-level overlap does not require `create_overlap_config`; `dispatch_async` +and `combine_async` can use their default stream/event behavior. Use +`create_overlap_config` when the caller wants explicit stream/event/SM settings +or block-level combine overlap. + +For block-level MLP/combine overlap, the config includes the combine-side wait +protocol and the device signal that an overlap-capable MLP backend must publish: + +```python +combine_overlap_config = moe_comm.create_overlap_config( + op="combine", + handle=handle, + level="block", +) +``` + +`op="dispatch", level="block"` is not part of the first version. Dispatch +overlap is operation-level only. + +`CommOverlapConfig` fields: + +| Field | Purpose | +|---|---| +| `op` | `"dispatch"` or `"combine"` | +| `level` | `"op"` or `"block"` | +| `stream` | Optional communication stream | +| `wait_event` | Optional event the communication op waits on before starting | +| `signal` | Device tensor written by MLP and waited on by combine for block overlap | +| `num_comm_sms` | Optional SM budget for communication | +| `block_m` | Rows per block for block overlap | +| `block_ready_value` | Signal value that marks one block as ready for combine | + +`DispatchHandle` should store the metadata needed to reverse dispatch: + +- source rank and source token index, +- top-k slot or equivalent routing metadata, +- top-k ids and routing weights, or stable references/copies, +- dispatch layout/range/count metadata, +- capacity, local expert placement, and launch parameters needed by kernels, +- optional cached metadata for repeated routing. + +## Dispatch inputs + +### `input` + +`input` is the original local token matrix. + +```python +input: torch.Tensor # [num_tokens, hidden] +``` + +Requirements: + +| Field | Requirement | +|---|---| +| Shape | `[T, H]`, token-major | +| Layout | contiguous row-major | +| Device | CUDA tensor | +| dtype | BF16, FP16, FP8, NVFP4, MXFP8, or another supported activation dtype | +| Ordering | original local token order; not expert sorted | + +The user should not expand `input` by top-k and should not convert it to +expert-major before calling `dispatch`. + +`dispatch` includes any metadata exchange needed before moving token payloads. +For normal/high-throughput modes this typically means computing send counts from +`topk_ids`, exchanging counts or layout information across ranks, choosing recv +slots, and then dispatching the activation payload. Users should not call a +separate metadata-exchange API in the simple path. + +### `topk_ids` + +```python +topk_ids: torch.Tensor # [T, K], int64 +``` + +`topk_ids[t, k]` is the global expert id selected for token `t` at top-k slot +`k`. Invalid slots may use `-1` if supported by the backend. + +### `weights` + +```python +weights: Optional[torch.Tensor] # [T, K], usually float32 +``` + +These are MoE routing weights, not quantization scales. They are used by +combine to reduce the `K` expert results for each token back to `[T, H]`. + +### `scales` + +`scales` contains activation quantization metadata for `input`. It should be +`None` for BF16/FP16 input. + +Examples: + +| Format | `input` | `scales.local` | `scales.global_scale` | +|---|---|---|---| +| BF16/FP16 | `[T, H]` | `None` | `None` | +| FP8 E4M3 | `[T, H]` FP8 | `[T, H / block_size]`, often block size 128 | usually `None` | +| NVFP4 | backend-defined packed/logical `[T, H]` | block scale tensor | optional global scale | +| MXFP8 | backend-defined `[T, H]` | micro-scale tensor, e.g. E8M0 blocks | optional/global if required | + +The API should not assume quantization scale is a scalar. For FP8 paths in +DeepEP/SGLang, scales are usually per token and per hidden block. + +### `output_buffer` + +Low-latency dispatch requires the caller to provide the receive token buffer: + +```python +output_buffer: torch.Tensor +``` + +For padded expert-major LL layout: + +```text +output_buffer: [num_local_experts, world_size * max_tokens_per_rank, hidden] +``` + +The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16. +For FP8 dispatch it is FP8 and the returned `DispatchOutput.scales` carries the +matching scale tensor. + +`output_buffer` is required for LL because the MLP runner often owns or reuses +workspace memory. `MoECommunicator` writes dispatch output into the provided +buffer instead of allocating it internally. + +## Dispatch output layout for MLP + +`dispatch` should return MLP-ready tokens. The MLP should not run another +token-major to expert-major permutation unless it uses a custom adapter. + +### Normal / high-throughput flat layout + +HT uses `DispatchLayout.FLAT`, a flat expert-major layout: + +```python +dispatch_out.tokens # [total_recv_tokens, H] +``` + +Rows are grouped by local expert id: + +```text +expert0 tokens +expert1 tokens +expert2 tokens +... +``` + +`dispatch_out.num_tokens_per_expert` is ordered by local expert id: + +```python +num_tokens_per_expert[i] = valid token count for local expert i +``` + +For flat layout, `expert_offsets` may be provided or derived by cumulative sum: + +```python +expert_offsets = cumsum([0] + num_tokens_per_expert) +tokens[expert_offsets[i] : expert_offsets[i + 1]] +``` + +This layout is efficient for Triton or grouped GEMM kernels because it avoids +padding. + +### Low-latency output layouts + +LL defaults to `DispatchLayout.EXPERT_MAJOR`, a padded expert-major tensor: + +```python +dispatch_out.tokens # [num_local_experts, max_slots_per_expert, H] +``` + +LL can also return `DispatchLayout.FLAT`, which is the same contiguous +local-expert-major storage viewed as 2D: + +```python +dispatch_out.tokens # [num_local_experts * max_slots_per_expert, H] +``` + +For expert `i`, only the first `num_tokens_per_expert[i]` slots are valid: + +```python +expert_major_tokens = dispatch_out.tokens.view(num_local_experts, max_slots_per_expert, H) +expert_major_tokens[i, :num_tokens_per_expert[i], :] +``` + +The remaining slots are padding or scratch space. The MLP output must keep the +same layout and slot order. + +### Scale output layout + +If `dispatch_out.scales` is not `None`, its local scale tensor should follow +the same packed/expert-major layout as `dispatch_out.tokens`, with the hidden +dimension replaced by the scale dimension. + +Examples: + +```text +flat tokens: HT [total_recv_tokens, H]; LL [num_local_experts * max_slots, H] +flat FP8 scales: HT [total_recv_tokens, H / 128]; LL [num_local_experts, max_slots, H / 128] + +expert-major tokens: [num_local_experts, max_slots, H] +expert-major scales: [num_local_experts, max_slots, H / 128] +``` + +## MLP contract + +The MLP consumes `dispatch_out`, not the original token-major input. + +For flat expert-major output: + +```python +expert_output = triton_mlp( + dispatch_out.tokens, + dispatch_out.num_tokens_per_expert, + dispatch_out.scales, +) +``` + +For padded expert-major output: + +```python +expert_output = expert_major_mlp( + dispatch_out.tokens, + dispatch_out.num_tokens_per_expert, + dispatch_out.scales, +) +``` + +The MLP must preserve the dispatch output layout and row/slot order. It may +apply expert-specific GEMMs, but it must not compact or reorder tokens unless it +also produces compatible metadata for combine. + +## Combine API + +```python +output = moe_comm.combine( + expert_output, + handle, + out=None, +) +``` + +`expert_output` must have the same physical layout and order as +`dispatch_out.tokens`. + +`combine` uses `handle` to: + +- map each expert output row/slot back to the original source rank and token, +- find the corresponding top-k slot, +- apply the routing weight, +- reduce all selected expert outputs for each token, +- return local output in original token-major order. + +The output shape is: + +```python +output # [T, H] +``` + +`combine` should not require users to pass `topk_ids`, `weights`, prefix +matrices, source indices, or layout ranges again. Those belong in `handle`. + +An optional `weights` override could be added later, but the default API should +use the weights captured by `dispatch`. + +## Communication/computation overlap + +The default API should be blocking and simple: + +```python +dispatch_out, handle = moe_comm.dispatch( + input, + topk_ids, + weights, + scales, + output_buffer=output_buffer, +) +expert_output = mlp(dispatch_out.tokens, dispatch_out.num_tokens_per_expert) +output = moe_comm.combine(expert_output, handle) +``` + +For overlap, expose two optional APIs rather than adding many flags to the +default path: + +| API | Purpose | +|---|---| +| `dispatch_async` / `combine_async` | Coarse-grained async launch and wait | +| `create_overlap_config(..., level="block")` | Fine-grained block notify between MLP down-GEMM and combine | + +### Coarse-grained overlap + +Coarse-grained overlap lets the caller launch communication on a communication +stream and wait later. + +```python +dispatch_overlap_config = moe_comm.create_overlap_config(op="dispatch") +dispatch_req = moe_comm.dispatch_async( + input, + topk_ids, + weights, + scales, + output_buffer=output_buffer, + overlap_config=dispatch_overlap_config, +) + +# Run unrelated work while dispatch metadata/payload communication is in flight. + +dispatch_out, handle = dispatch_req.wait() +expert_output = mlp(dispatch_out.tokens, dispatch_out.num_tokens_per_expert) + +combine_overlap_config = moe_comm.create_overlap_config(op="combine", handle=handle) +combine_req = moe_comm.combine_async( + expert_output, + handle, + overlap_config=combine_overlap_config, +) + +# Run unrelated work while combine is in flight. + +output = combine_req.wait() +``` + +This is similar to the event/hook style used by DeepEP and SGLang. The request +object should own any stream event or receive hook required by the backend. + +### Fine-grained MLP/combine overlap + +Fine-grained overlap sends combine data as soon as the MLP produces a block. +This requires a device-side notify/signal from the MLP backend to the combine +kernel. + +```python +combine_overlap_config = moe_comm.create_overlap_config( + op="combine", + handle=handle, + level="block", +) + +# User must adapt the MLP backend/adapter to consume this config and notify +# combine as blocks become ready. +config = combine_overlap_config +expert_output = mlp( + dispatch_out.tokens, + dispatch_out.num_tokens_per_expert, + config=config, +) + +combine_req = moe_comm.combine_async( + expert_output, + handle, + overlap_config=combine_overlap_config, +) + +output = combine_req.wait() +``` + +The overlap config is not routing metadata. It only tells combine when a +region of `expert_output` is ready to read. The routing/source mapping still +comes from `handle`. + +The MLP backend must follow these rules when using notify: + +- write `expert_output` in the same row/slot order as `dispatch_out.tokens`, +- publish data before signaling readiness, +- signal at the block granularity defined by `overlap_config`, +- use the signal value/protocol provided by `overlap_config`. + +If the MLP backend does not support notify, it can still use the blocking API or +coarse-grained `combine_async` after the full `expert_output` tensor is ready. + +This must be a joint contract between the dispatcher and the MLP runner. The +dispatcher can provide the signal buffer and combine-side wait protocol, but it +cannot infer readiness by itself. The MLP runner must write the signal after it +finishes the corresponding output region. + +SGLang follows this model for its DeepEP low-latency path. It computes overlap +arguments after dispatch, passes combine-side arguments to the DeepEP dispatcher, +and passes down-GEMM arguments to the MoE runner. Backend support is selective: + +- DeepGEMM FP8 masked down-GEMM can return block metadata such as `block_m` and + `block_ready_value` and signal combine readiness. +- FlashInfer CuteDSL can receive down-GEMM signal/start-event arguments. +- Some paths, such as BF16 masked DeepGEMM and generic Triton runners, do not + support this block overlap protocol. + +Therefore, the API should expose overlap as an optional capability advertised by +the MLP backend, not as a guaranteed feature of every `combine_async` call. + +## Internal metadata exchange + +Normal/high-throughput dispatch usually needs a metadata phase before payload +movement: + +```text +topk_ids + -> compute send counts per rank/expert + -> exchange counts or layout metadata + -> compute recv slots and local expert counts + -> dispatch token payload +``` + +Low-latency modes may use fixed-capacity buffers and device-side counters, but +they still generate metadata such as source info, layout ranges, and valid +counts. + +These details should remain internal. The user-facing API should only expose +MLP-relevant layout information through `DispatchOutput` and combine-relevant +metadata through `DispatchHandle`. + +## Example + +```python +recv, handle = moe_comm.dispatch( + input=hidden_states, # [T, H] + topk_ids=topk_ids, # [T, K] + weights=topk_weights, # [T, K] + scales=None, # BF16 path + output_buffer=recv_buffer, +) + +expert_output = triton_grouped_mlp( + recv.tokens, + recv.num_tokens_per_expert, +) + +output = moe_comm.combine(expert_output, handle) +``` + +Quantized path: + +```python +recv, handle = moe_comm.dispatch( + input=x_fp8, + topk_ids=topk_ids, + weights=topk_weights, + scales=QuantScales( + local=x_scales, + format="fp8_e4m3", + block_size=128, + ), + output_buffer=recv_buffer, +) + +expert_output = fp8_grouped_mlp( + recv.tokens, + recv.scales, + recv.num_tokens_per_expert, +) + +output = moe_comm.combine(expert_output, handle) +``` diff --git a/python/mscclpp/ext/ep/__init__.py b/python/mscclpp/ext/ep/__init__.py index 1ed94491..a551b9c5 100644 --- a/python/mscclpp/ext/ep/__init__.py +++ b/python/mscclpp/ext/ep/__init__.py @@ -2,12 +2,28 @@ # Licensed under the MIT License. """MSCCL++ Expert-Parallel (MoE dispatch/combine) extension. -See ``src/ext/ep/README.md`` in the repository for migration status. The -``Buffer`` class mirrors :class:`deep_ep.Buffer` and supports intranode -(NVLink-only) dispatch/combine as well as internode HT and low-latency -paths. +See ``src/ext/ep/README.md`` in the repository for migration status. +``MoECommunicator`` is the high-level public API. """ -from .buffer import Buffer, Config, EventHandle # noqa: F401 +from .communicator import ( # noqa: F401 + CommOverlapConfig, + DispatchHandle, + DispatchOutput, + DispatchLayout, + MoEMode, + MoECommunicator, + MoECommunicatorConfig, + QuantScales, +) -__all__ = ["Buffer", "Config", "EventHandle"] +__all__ = [ + "CommOverlapConfig", + "DispatchHandle", + "DispatchOutput", + "DispatchLayout", + "MoEMode", + "MoECommunicator", + "MoECommunicatorConfig", + "QuantScales", +] diff --git a/python/mscclpp/ext/ep/buffer.py b/python/mscclpp/ext/ep/buffer.py deleted file mode 100644 index 1c2f9604..00000000 --- a/python/mscclpp/ext/ep/buffer.py +++ /dev/null @@ -1,193 +0,0 @@ -# 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. -"""Python frontend for the MSCCL++ Expert-Parallel extension. - -This is a thin wrapper around the pybind11 extension ``mscclpp_ep_cpp``. -The shape of :class:`Buffer` mirrors :class:`deep_ep.Buffer` so existing -DeepEP users can port with minimal changes. - -Current status (see ``src/ext/ep/README.md``): - -* Intranode (NVLink-only) dispatch and combine: ported and validated on - one node with 8 GPUs. -* ``get_dispatch_layout``: ported. -* Internode HT (MSCCL++ PortChannel + MemoryChannel) dispatch and combine: - ported and validated on 2 nodes x 8 H100 GPUs with - ``test/python/ext/ep/test_internode_multirank.py``. -* Internode low-latency kernels (NVSHMEM/IBGDA -> MSCCL++ PortChannel): - ported and validated on 2 nodes x 8 H100 GPUs with - ``test/python/ext/ep/test_low_latency_multirank.py``. -""" - -from __future__ import annotations - -import os -from typing import List, Optional, Tuple - -import torch -import torch.distributed as dist - -try: - import mscclpp_ep_cpp as _cpp # type: ignore[import-not-found] -except ImportError as exc: # pragma: no cover - raise ImportError( - "mscclpp_ep_cpp is not available. Build mscclpp with " - "-DMSCCLPP_BUILD_EXT_EP=ON (and ensure PyTorch's CMake prefix is on " - "CMAKE_PREFIX_PATH) or install via `pip install` after the build." - ) from exc - -Config = _cpp.Config -EventHandle = _cpp.EventHandle - - -class Buffer: - """Core expert-parallel (EP) communication buffer. - - Parameters - ---------- - group: - The ``torch.distributed`` process group. Used only for out-of-band - exchange of IPC handles and the MSCCL++ unique id. - num_nvl_bytes: - Size of the NVLink-accessible scratch buffer (shared via CUDA IPC). - num_rdma_bytes: - Size of the RDMA scratch buffer. Required (>0) for internode HT and - low-latency modes. - low_latency_mode: - Enable the low-latency dispatch/combine path. This mode uses only - the RDMA buffer (``num_rdma_bytes``) and drives every peer through - MSCCL++ ``PortChannel``; consequently, it works cross-node with any - topology but is still pending H100 hardware validation. - num_qps_per_rank: - Ignored for intranode mode. - """ - - #: Default number of SMs reserved for comms kernels. Matches DeepEP. - num_sms: int = 20 - - def __init__( - self, - group: dist.ProcessGroup, - num_nvl_bytes: int = 0, - num_rdma_bytes: int = 0, - low_latency_mode: bool = False, - num_qps_per_rank: int = 12, - ) -> None: - self.rank: int = group.rank() - self.group_size: int = group.size() - self.group = group - self.num_nvl_bytes = num_nvl_bytes - self.num_rdma_bytes = num_rdma_bytes - self.low_latency_mode = low_latency_mode - self.num_qps_per_rank = num_qps_per_rank - - self.runtime = _cpp.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode) - - # Exchange device IDs + IPC handles + (for RDMA) the MSCCL++ unique id. - device_ids: List[Optional[int]] = [None] * self.group_size - local_device_id = self.runtime.get_local_device_id() - dist.all_gather_object(device_ids, local_device_id, group) - - ipc_handles: List[Optional[bytes]] = [None] * self.group_size - local_ipc_handle = self.runtime.get_local_ipc_handle() - dist.all_gather_object(ipc_handles, local_ipc_handle, group) - - root_unique_id: Optional[bytes] = None - # MSCCL++ requires a bootstrapped Communicator even for pure-NVLink - # setups because `Buffer::sync()` uses `communicator->connect(ipc)` - # to build MemoryChannels. We always exchange a unique id. - if num_qps_per_rank <= 0: - raise ValueError("num_qps_per_rank must be > 0") - - if self.rank == 0: - root_unique_id = self.runtime.create_unique_id() - broadcast_list = [root_unique_id] - dist.broadcast_object_list(broadcast_list, src=0, group=group) - root_unique_id = broadcast_list[0] - assert root_unique_id is not None - self.runtime.connect(root_unique_id) - - # sync() expects Sequence[bytearray | None] / bytearray | None. - 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)) - - # ------------------------------------------------------------------ - # Sanity helpers - # ------------------------------------------------------------------ - - def is_available(self) -> bool: - return self.runtime.is_available() - - 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_) - - # ------------------------------------------------------------------ - # Layout / dispatch / combine (thin pass-through wrappers). - # Signatures mirror deep_ep.Buffer so existing test harnesses can reuse. - # ------------------------------------------------------------------ - - def get_dispatch_layout( - self, - topk_idx: torch.Tensor, - num_experts: int, - previous_event: Optional[EventHandle] = None, - async_finish: bool = False, - allocate_on_comm_stream: bool = False, - ): - return self.runtime.get_dispatch_layout( - topk_idx, num_experts, previous_event, async_finish, allocate_on_comm_stream - ) - - def intranode_dispatch(self, *args, **kwargs): - return self.runtime.intranode_dispatch(*args, **kwargs) - - def intranode_combine(self, *args, **kwargs): - return self.runtime.intranode_combine(*args, **kwargs) - - def internode_dispatch(self, *args, **kwargs): - return self.runtime.internode_dispatch(*args, **kwargs) - - def internode_combine(self, *args, **kwargs): - return self.runtime.internode_combine(*args, **kwargs) - - def clean_low_latency_buffer(self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> None: - self.runtime.clean_low_latency_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts) - - def low_latency_dispatch(self, *args, **kwargs): - return self.runtime.low_latency_dispatch(*args, **kwargs) - - def low_latency_combine(self, *args, **kwargs): - return self.runtime.low_latency_combine(*args, **kwargs) - - def get_next_low_latency_combine_buffer(self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int): - return self.runtime.get_next_low_latency_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts) - - def get_local_buffer_tensor( - self, dtype: torch.dtype, offset: int = 0, use_rdma_buffer: bool = False - ) -> torch.Tensor: - return self.runtime.get_local_buffer_tensor(dtype, offset, use_rdma_buffer) - - # ------------------------------------------------------------------ - # Static helpers - # ------------------------------------------------------------------ - - @staticmethod - def get_low_latency_rdma_size_hint( - num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int - ) -> int: - return _cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts) diff --git a/python/mscclpp/ext/ep/communicator.py b/python/mscclpp/ext/ep/communicator.py new file mode 100644 index 00000000..36648f12 --- /dev/null +++ b/python/mscclpp/ext/ep/communicator.py @@ -0,0 +1,536 @@ +# 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. +"""Python frontend for the MSCCL++ Expert-Parallel extension. + +This is a thin wrapper around the nanobind extension ``mscclpp_ep_cpp``. +``MoECommunicator`` is the high-level API. ``_MoERuntime`` is a private +low-latency runtime wrapper used internally by the high-level API. + +Current status (see ``src/ext/ep/README.md``): + +* Intranode (NVLink-only) dispatch and combine: ported and validated on + one node with 8 GPUs. +* ``get_dispatch_layout``: ported. +* Internode HT (MSCCL++ PortChannel + MemoryChannel) dispatch and combine: + ported and validated on 2 nodes x 8 H100 GPUs with + ``test/python/ext/ep/test_internode_multirank.py``. +* Low-latency kernels (RDMA + CUDA IPC paths): + ported and validated on intra-node and 2 nodes x 8 H100 GPUs with + ``test/python/ext/ep/test_low_latency_multirank.py``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch +from mscclpp._core import CommGroup + +try: + import mscclpp_ep_cpp as _cpp # type: ignore[import-not-found] +except ImportError as exc: # pragma: no cover + raise ImportError( + "mscclpp_ep_cpp is not available. Build mscclpp with -DMSCCLPP_BUILD_EXT_EP=ON " + "or install with `pip install .[ep]`." + ) from exc + + +DispatchLayout = _cpp.DispatchLayout +MoEMode = _cpp.MoEMode + + +@dataclass +class MoECommunicatorConfig: + """Configuration for the high-level MoE dispatch/combine API.""" + + comm: Optional[CommGroup] = None + device: Optional[torch.device | int] = None + num_experts: int = 0 + num_local_experts: Optional[int] = None + local_expert_start: Optional[int] = None + hidden_size: int = 0 + topk: int = 0 + max_tokens_per_rank: int = 0 + max_recv_tokens_per_rank: Optional[int] = None + mode: MoEMode = MoEMode.LOW_LATENCY + output_layout: Optional[DispatchLayout] = None + input_dtype: Optional[torch.dtype] = None + quant_format: Optional[str] = None + num_rdma_qps_per_rank: int = 12 + num_sms: int = 20 + enable_overlap: bool = False + + +@dataclass +class QuantScales: + local: Optional[torch.Tensor] = None + global_scale: Optional[torch.Tensor] = None + format: Optional[str] = None + block_size: Optional[int] = None + + +@dataclass +class DispatchOutput: + tokens: torch.Tensor + scales: Optional[QuantScales] + num_tokens_per_expert: torch.Tensor | list[int] + expert_offsets: Optional[torch.Tensor] = None + layout: DispatchLayout = DispatchLayout.FLAT + + +@dataclass +class DispatchHandle: + """Opaque dispatch metadata consumed by :meth:`MoECommunicator.combine`.""" + + topk_ids: torch.Tensor + weights: torch.Tensor + src_info: torch.Tensor + layout_range: torch.Tensor + num_max_dispatch_tokens_per_rank: int + num_experts: int + num_tokens: int + hidden_size: int + num_local_experts: int + local_expert_start: int + layout: DispatchLayout + output_scales: Optional[QuantScales] = None + + +@dataclass +class CommOverlapConfig: + op: str + level: str = "op" + stream: Optional[torch.cuda.Stream] = None + wait_event: Optional[torch.cuda.Event] = None + signal: Optional[torch.Tensor] = None + num_comm_sms: Optional[int] = None + block_m: Optional[int] = None + block_ready_value: Optional[int] = None + + +class _MoERuntime: + """Private low-level MoE communication runtime wrapper. + + Parameters + ---------- + comm: + The :class:`mscclpp.CommGroup`. Used only for out-of-band + exchange of IPC handles and the MSCCL++ unique id. + num_nvl_bytes: + Size of the NVLink-accessible scratch buffer. Reserved for archived HT mode. + num_rdma_bytes: + Size of the LL RDMA scratch buffer. + mode: + Runtime mode selector. ``MoEMode.LOW_LATENCY`` is active; + ``MoEMode.HIGH_THROUGHPUT`` is archived and not compiled. + num_qps_per_rank: + RDMA QPs per peer rank. + """ + + #: Default number of SMs reserved for comms kernels. Matches DeepEP. + num_sms: int = 20 + + def __init__( + self, + comm: CommGroup, + num_nvl_bytes: int = 0, + num_rdma_bytes: int = 0, + mode: MoEMode = MoEMode.LOW_LATENCY, + num_qps_per_rank: int = 12, + ) -> None: + if not isinstance(mode, MoEMode): + raise TypeError("mode must be a MoEMode") + self.mode = mode + if self.mode != MoEMode.LOW_LATENCY: + raise NotImplementedError("mode='ht' is archived under src/ext/ep/ht and is not compiled") + + 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._cpp_runtime = _cpp.MoERuntime(comm.communicator, num_nvl_bytes, num_rdma_bytes, mode) + + if num_qps_per_rank <= 0: + raise ValueError("num_qps_per_rank must be > 0") + + # ------------------------------------------------------------------ + # Sanity helpers + # ------------------------------------------------------------------ + + def is_available(self) -> bool: + return self._cpp_runtime.is_available() + + def is_internode_available(self) -> bool: + return self._cpp_runtime.is_internode_available() + + def get_local_device_id(self) -> int: + return self._cpp_runtime.get_local_device_id() + + def get_num_rdma_ranks(self) -> int: + return self._cpp_runtime.get_num_rdma_ranks() + + def get_rdma_rank(self) -> int: + return self._cpp_runtime.get_rdma_rank() + + def get_root_rdma_rank(self, global_: bool) -> int: + return self._cpp_runtime.get_root_rdma_rank(global_) + + +class MoECommunicator: + """High-level MoE communicator API for dispatch/combine. + + The first implementation supports the low-latency backend. + """ + + def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> None: + if config is not None and kwargs: + raise ValueError("Pass either MoECommunicatorConfig or keyword arguments, not both") + if config is None: + if "group" in kwargs and "comm" not in kwargs: + kwargs["comm"] = kwargs.pop("group") + config = MoECommunicatorConfig(**kwargs) + + if config.device is not None: + torch.cuda.set_device(config.device) + + comm = config.comm + if comm is None: + raise ValueError("MoECommunicator requires an mscclpp.CommGroup") + + self.comm = comm + self.rank: int = comm.my_rank + self.world_size: int = comm.nranks + self.local_rank: int = torch.cuda.current_device() + self.device = torch.device("cuda", self.local_rank) + + if not isinstance(config.mode, MoEMode): + raise TypeError("MoECommunicatorConfig.mode must be a MoEMode") + self.mode = config.mode + if self.mode != MoEMode.LOW_LATENCY: + raise NotImplementedError("mode='ht' is archived under src/ext/ep/ht and is not compiled") + + self.output_layout = _resolve_output_layout(config.output_layout, self.mode) + + self.num_experts = config.num_experts + self.hidden_size = config.hidden_size + self.topk = config.topk + self.max_tokens_per_rank = config.max_tokens_per_rank + if self.num_experts <= 0 or self.hidden_size <= 0 or self.topk <= 0 or self.max_tokens_per_rank <= 0: + raise ValueError("num_experts, hidden_size, topk, and max_tokens_per_rank must be positive") + if self.num_experts % self.world_size != 0: + raise ValueError("low-latency mode requires num_experts divisible by world_size") + + self.num_local_experts = config.num_local_experts + if self.num_local_experts is None: + self.num_local_experts = self.num_experts // self.world_size + if self.num_local_experts * self.world_size != self.num_experts: + raise NotImplementedError("only even contiguous expert placement is currently supported") + + self.local_expert_start = config.local_expert_start + if self.local_expert_start is None: + self.local_expert_start = self.rank * self.num_local_experts + + if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank): + raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity") + if config.input_dtype not in (None, torch.bfloat16): + raise NotImplementedError("low-latency dispatch currently supports BF16 input only") + + self.quant_format = _normalize_quant_format(config.quant_format) + if self.quant_format not in (None, "fp8_e4m3"): + raise NotImplementedError(f"unsupported low-latency quant_format: {config.quant_format}") + self.dispatch_requires_quantization = self.quant_format is not None + + num_nvl_bytes = 0 + num_rdma_bytes = _get_low_latency_rdma_size_hint( + self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts + ) + + self.enable_overlap = config.enable_overlap + self.num_sms = config.num_sms + self._dispatch_scales: Optional[torch.Tensor] = None + self._dispatch_src_info: Optional[torch.Tensor] = None + self._dispatch_layout_range: Optional[torch.Tensor] = None + self._dispatch_count: Optional[torch.Tensor] = None + + self._runtime = _MoERuntime( + comm, + num_nvl_bytes=num_nvl_bytes, + num_rdma_bytes=num_rdma_bytes, + mode=self.mode, + num_qps_per_rank=config.num_rdma_qps_per_rank, + ) + + def is_available(self) -> bool: + return self._runtime.is_available() + + def is_internode_available(self) -> bool: + return self._runtime.is_internode_available() + + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor] = None, + scales: Optional[QuantScales] = None, + *, + output_buffer: torch.Tensor, + stream: Optional[torch.cuda.Stream] = None, + ) -> tuple[DispatchOutput, DispatchHandle]: + self._validate_dispatch_inputs(input, topk_ids, weights, scales, output_buffer) + if weights is None: + weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) + + output_tensors = self._get_dispatch_output_tensors(output_buffer) + output_buffer, packed_scales, src_info, layout_range, num_tokens_per_expert = output_tensors + self._runtime._cpp_runtime.dispatch( + input.data_ptr(), + topk_ids.data_ptr(), + output_buffer.data_ptr(), + 0 if packed_scales is None else packed_scales.data_ptr(), + src_info.data_ptr(), + layout_range.data_ptr(), + num_tokens_per_expert.data_ptr(), + input.size(0), + self.hidden_size, + self.topk, + self.max_tokens_per_rank, + self.num_experts, + self.dispatch_requires_quantization, + self.output_layout, + _cuda_stream_ptr(stream), + ) + output_scales = None + if packed_scales is not None: + output_scales = QuantScales(local=packed_scales, format="fp8_e4m3", block_size=128) + + dispatch_out = DispatchOutput( + tokens=output_buffer, + scales=output_scales, + num_tokens_per_expert=num_tokens_per_expert, + expert_offsets=None, + layout=self.output_layout, + ) + handle = DispatchHandle( + topk_ids=topk_ids, + weights=weights, + src_info=src_info, + layout_range=layout_range, + num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank, + num_experts=self.num_experts, + num_tokens=input.size(0), + hidden_size=self.hidden_size, + num_local_experts=self.num_local_experts, + local_expert_start=self.local_expert_start, + layout=self.output_layout, + output_scales=output_scales, + ) + return dispatch_out, handle + + def _get_dispatch_output_tensors(self, output_buffer: torch.Tensor) -> tuple[ + torch.Tensor, + Optional[torch.Tensor], + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + device = output_buffer.device + slots_per_expert = self.world_size * self.max_tokens_per_rank + if self._dispatch_src_info is None or self._dispatch_src_info.device != device: + self._dispatch_src_info = torch.empty( + (self.num_local_experts, slots_per_expert), + dtype=torch.int32, + device=device, + ) + self._dispatch_layout_range = torch.empty( + (self.num_local_experts, self.world_size), + dtype=torch.int64, + device=device, + ) + self._dispatch_count = torch.empty( + (self.num_local_experts,), + dtype=torch.int32, + device=device, + ) + self._dispatch_scales = None + if self.dispatch_requires_quantization: + num_scales = self.hidden_size // 128 + scales_storage = torch.empty( + (self.num_local_experts, num_scales, slots_per_expert), + dtype=torch.float32, + device=device, + ) + self._dispatch_scales = scales_storage.transpose(1, 2) + + assert self._dispatch_src_info is not None + assert self._dispatch_layout_range is not None + assert self._dispatch_count is not None + return ( + output_buffer, + self._dispatch_scales, + self._dispatch_src_info, + self._dispatch_layout_range, + self._dispatch_count, + ) + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor] = None, + stream: Optional[torch.cuda.Stream] = None, + ) -> torch.Tensor: + self._validate_combine_inputs(expert_output, handle, out) + combine_requires_dequantization = _requires_dequantization(expert_output) + x_scales = None + if combine_requires_dequantization: + if handle.output_scales is None or handle.output_scales.local is None: + raise ValueError("FP8 expert_output requires scales captured in the dispatch handle") + x_scales = handle.output_scales.local + + if out is None: + out = torch.empty((handle.num_tokens, self.hidden_size), dtype=torch.bfloat16, device=expert_output.device) + self._runtime._cpp_runtime.combine( + expert_output.data_ptr(), + 0 if x_scales is None else x_scales.data_ptr(), + handle.topk_ids.data_ptr(), + handle.weights.data_ptr(), + handle.src_info.data_ptr(), + handle.layout_range.data_ptr(), + out.data_ptr(), + handle.num_tokens, + self.hidden_size, + handle.weights.size(1), + handle.num_max_dispatch_tokens_per_rank, + handle.num_experts, + combine_requires_dequantization, + _cuda_stream_ptr(stream), + ) + return out + + def dispatch_async(self, *args, **kwargs): + raise NotImplementedError("dispatch_async is not implemented for MoECommunicator yet") + + def combine_async(self, *args, **kwargs): + raise NotImplementedError("combine_async is not implemented for MoECommunicator yet") + + def create_overlap_config( + self, + op: str, + *, + handle: Optional[DispatchHandle] = None, + level: str = "op", + ) -> CommOverlapConfig: + if op not in ("dispatch", "combine"): + raise ValueError("op must be 'dispatch' or 'combine'") + if level != "op": + raise NotImplementedError("block-level overlap is not implemented yet") + if op == "combine" and handle is None: + raise ValueError("combine overlap config requires a DispatchHandle") + return CommOverlapConfig(op=op, level=level) + + def _validate_dispatch_inputs( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + scales: Optional[QuantScales], + output_buffer: torch.Tensor, + ) -> None: + if output_buffer is None: + raise ValueError("output_buffer is required for low-latency dispatch") + if scales is not None and (scales.local is not None or scales.global_scale is not None): + raise NotImplementedError("low-latency dispatch does not support quantized input scales yet") + if input.dim() != 2 or not input.is_contiguous(): + raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor") + if input.device.type != "cuda" or input.dtype != torch.bfloat16: + raise ValueError("low-latency dispatch input must be a CUDA BF16 tensor") + if input.size(1) != self.hidden_size: + raise ValueError(f"input hidden size {input.size(1)} does not match configured {self.hidden_size}") + if input.size(0) > self.max_tokens_per_rank: + raise ValueError("input token count exceeds max_tokens_per_rank") + if topk_ids.dim() != 2 or not topk_ids.is_contiguous(): + raise ValueError("topk_ids must be a contiguous [num_tokens, topk] tensor") + if topk_ids.device != input.device or topk_ids.dtype != torch.int64: + raise ValueError("topk_ids must be an int64 CUDA tensor on the same device as input") + if topk_ids.shape != (input.size(0), self.topk): + raise ValueError("topk_ids shape must match [input.size(0), configured topk]") + if weights is not None: + if weights.dim() != 2 or not weights.is_contiguous(): + raise ValueError("weights must be a contiguous [num_tokens, topk] tensor") + if weights.device != input.device or weights.dtype != torch.float32: + raise ValueError("weights must be a float32 CUDA tensor on the same device as input") + if weights.shape != topk_ids.shape: + raise ValueError("weights shape must match topk_ids") + expected_dtype = torch.float8_e4m3fn if self.dispatch_requires_quantization else torch.bfloat16 + slots_per_expert = self.world_size * self.max_tokens_per_rank + if self.output_layout == DispatchLayout.EXPERT_MAJOR: + expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size) + else: + expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size) + if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous(): + raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor") + if output_buffer.device != input.device or output_buffer.dtype != expected_dtype: + raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input") + if tuple(output_buffer.shape) != expected_shape: + raise ValueError(f"output_buffer shape must be {expected_shape}") + + def _validate_combine_inputs( + self, expert_output: torch.Tensor, handle: DispatchHandle, out: Optional[torch.Tensor] + ) -> None: + if handle.num_experts != self.num_experts or handle.hidden_size != self.hidden_size: + raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration") + slots_per_expert = self.world_size * self.max_tokens_per_rank + if handle.layout == DispatchLayout.EXPERT_MAJOR: + expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size) + else: + expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size) + if expert_output.dim() != len(expected_shape) or not expert_output.is_contiguous(): + raise ValueError("expert_output must keep dispatch output's contiguous layout") + if tuple(expert_output.shape) != expected_shape: + raise ValueError(f"expert_output shape must be {expected_shape}") + if expert_output.dtype not in (torch.bfloat16, getattr(torch, "float8_e4m3fn", None)): + raise ValueError("expert_output must be BF16 or FP8 E4M3") + if out is not None: + expected_out_shape = (handle.num_tokens, self.hidden_size) + if tuple(out.shape) != expected_out_shape or out.dtype != torch.bfloat16 or not out.is_contiguous(): + raise ValueError(f"out must be a contiguous BF16 tensor with shape {expected_out_shape}") + + +def _resolve_output_layout(layout: Optional[DispatchLayout], mode: MoEMode) -> DispatchLayout: + if layout is None: + return DispatchLayout.EXPERT_MAJOR if mode == MoEMode.LOW_LATENCY else DispatchLayout.FLAT + if not isinstance(layout, DispatchLayout): + raise TypeError("MoECommunicatorConfig.output_layout must be a DispatchLayout") + return layout + + +def _cuda_stream_ptr(stream: Optional[torch.cuda.Stream]) -> int: + return (stream if stream is not None else torch.cuda.current_stream()).cuda_stream + + +def _normalize_quant_format(fmt: Optional[str]) -> Optional[str]: + if fmt is None: + return None + normalized = fmt.lower().replace("-", "_") + if normalized in ("fp8", "fp8_e4m3", "f8e4m3", "float8_e4m3fn"): + return "fp8_e4m3" + return normalized + + +def _requires_dequantization(tensor: torch.Tensor) -> bool: + fp8_dtype = getattr(torch, "float8_e4m3fn", None) + return fp8_dtype is not None and tensor.dtype == fp8_dtype + + +def _get_low_latency_rdma_size_hint( + num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int +) -> int: + return _cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts) diff --git a/python/test/test_gpu_buffer_pool_nvls_zero.py b/python/test/test_gpu_buffer_pool_nvls_zero.py new file mode 100644 index 00000000..51a3e34d --- /dev/null +++ b/python/test/test_gpu_buffer_pool_nvls_zero.py @@ -0,0 +1,135 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import socket + +import cupy as cp +import pytest + +from mscclpp import CommGroup, DataType, RawGpuBuffer, ReduceOp, GpuBufferPool, is_nvls_supported +from mscclpp.ext import AlgorithmCollectionBuilder +from mscclpp_benchmark.gpu import capture_graph +from .mscclpp_mpi import MpiGroup, parametrize_mpi_groups, mpi_group # noqa: F401 + + +def _same_host(comm) -> bool: + hostnames = comm.allgather(socket.gethostname()) + return len(set(hostnames)) == 1 + + +def _build_nvls_zero_algorithm(mpi_group: MpiGroup): + comm_group = CommGroup(mpi_group.comm) + scratch = RawGpuBuffer(1 << 27) + AlgorithmCollectionBuilder.reset() + builder = AlgorithmCollectionBuilder() + algorithms = builder.build_default_algorithms( + scratch_buffer=scratch.data(), + scratch_buffer_size=scratch.bytes(), + rank=comm_group.my_rank, + ) + for algorithm in algorithms: + if algorithm.name == "default_allreduce_nvls_zero_copy": + return comm_group, algorithm, scratch + pytest.skip("default_allreduce_nvls_zero_copy is not available") + + +def _torch_tensor_from_pool_buffer(torch, buffer, nelems: int): + return torch.utils.dlpack.from_dlpack(buffer.to_dlpack(data_type=str(torch.float32), shape=[nelems])) + + +def _run_nvls_zero_copy(algorithm, comm_group, buffer, stream) -> None: + ret = algorithm.execute( + comm=comm_group.communicator, + input_buffer=buffer.data(), + output_buffer=buffer.data(), + input_size=buffer.bytes(), + output_size=buffer.bytes(), + dtype=DataType.float32, + op=ReduceOp.SUM, + stream=stream.ptr, + nblocks=0, + nthreads_per_block=0, + symmetric_memory=True, + accum_dtype=DataType.float32, + ) + assert ret == 0 + + +@parametrize_mpi_groups(2, 4, 8) +def test_gpu_buffer_pool_allreduce_nvls_zero_copy_timing(mpi_group: MpiGroup): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("Torch CUDA is not available") + if not is_nvls_supported(): + pytest.skip("NVLS is not supported") + if not _same_host(mpi_group.comm): + pytest.skip("NVLS zero-copy test requires all ranks on the same host") + + torch.cuda.set_device(cp.cuda.Device().id) + comm_group, algorithm, scratch = _build_nvls_zero_algorithm(mpi_group) + stream = cp.cuda.Stream(non_blocking=True) + + message_sizes = (256 * 1024, 1024 * 1024) + element_size = torch.empty((), dtype=torch.float32, device="cuda").element_size() + n_warmup = 3 + n_iters = 10 + pool = GpuBufferPool(sum(nbytes + 4096 for nbytes in message_sizes)) + expected = float(comm_group.nranks * (comm_group.nranks + 1) // 2) + live_tensors = [] + graphs = [] + + try: + for nbytes in message_sizes: + nelems = nbytes // element_size + buffer = pool.allocate(nbytes, alignment=4096) + tensor = _torch_tensor_from_pool_buffer(torch, buffer, nelems) + tensor.fill_(float(comm_group.my_rank + 1)) + torch.cuda.synchronize() + mpi_group.comm.barrier() + + _run_nvls_zero_copy(algorithm, comm_group, buffer, stream) + stream.synchronize() + assert torch.allclose(tensor, torch.full_like(tensor, expected)) + + tensor.fill_(float(comm_group.my_rank + 1)) + torch.cuda.synchronize() + mpi_group.comm.barrier() + + graph = capture_graph(stream, lambda: _run_nvls_zero_copy(algorithm, comm_group, buffer, stream)) + graphs.append(graph) + graph.launch(stream) + stream.synchronize() + assert torch.allclose(tensor, torch.full_like(tensor, expected)) + + for _ in range(n_warmup): + graph.launch(stream) + stream.synchronize() + mpi_group.comm.barrier() + + start = cp.cuda.Event() + end = cp.cuda.Event() + start.record(stream) + for _ in range(n_iters): + graph.launch(stream) + end.record(stream) + end.synchronize() + mpi_group.comm.barrier() + + elapsed_us = cp.cuda.get_elapsed_time(start, end) * 1000.0 / n_iters + all_elapsed_us = mpi_group.comm.allgather(elapsed_us) + if comm_group.my_rank == 0: + avg_us = max(all_elapsed_us) + print( + f"default_allreduce_nvls_zero_copy graph with GpuBufferPool: " + f"nranks={comm_group.nranks}, nbytes={nbytes}, avg={avg_us:.2f} us" + ) + live_tensors.append(tensor) + del buffer + + finally: + for graph in graphs: + graph.close() + live_tensors.clear() + torch.cuda.synchronize() + AlgorithmCollectionBuilder.reset() + del scratch diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index 3e60bdb2..cfa58e8f 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -1,48 +1,39 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -# Builds `mscclpp_ep_cpp`, a pybind11 + PyTorch C++ extension that exposes the -# EP (Mixture-of-Experts dispatch/combine) Buffer to Python. This module is -# separate from the nanobind-based `_mscclpp` because EP carries `torch::Tensor` -# through its ABI (ported verbatim from DeepEP). -# -# Requires: PyTorch with its CMake integration available on CMAKE_PREFIX_PATH. -# Easiest invocation: -# cmake -S . -B build \ -# -DMSCCLPP_BUILD_EXT_EP=ON \ -# -DCMAKE_PREFIX_PATH="$(python -c 'import torch;print(torch.utils.cmake_prefix_path)')" +# Builds `mscclpp_ep_cpp`, a nanobind extension that exposes the EP +# (Mixture-of-Experts dispatch/combine) runtime to Python. find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) -find_package(Torch QUIET) -if(NOT Torch_FOUND) - message(WARNING - "MSCCLPP_BUILD_EXT_EP=ON but PyTorch CMake package was not found. " - "Set CMAKE_PREFIX_PATH to `$(python -c 'import torch;print(torch.utils.cmake_prefix_path)')` " - "or disable MSCCLPP_BUILD_EXT_EP. Skipping the EP extension.") - return() +include(FetchContent) +if(NOT TARGET nanobind-static AND NOT TARGET nanobind) + FetchContent_Declare(nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git GIT_TAG v1.9.2) + FetchContent_MakeAvailable(nanobind) endif() -find_package(pybind11 QUIET CONFIG) -if(NOT pybind11_FOUND) - # PyTorch ships a bundled pybind11 we can fall back on. - get_target_property(_torch_include_dirs torch INTERFACE_INCLUDE_DIRECTORIES) - foreach(d ${_torch_include_dirs}) - if(EXISTS "${d}/pybind11/pybind11.h") - set(PYBIND11_INCLUDE_DIR "${d}") - break() +set(_mscclpp_ep_gpu_archs "") +if(MSCCLPP_USE_CUDA) + foreach(_arch IN LISTS MSCCLPP_GPU_ARCHS) + if(_arch STREQUAL "native") + list(APPEND _mscclpp_ep_gpu_archs "${_arch}") + else() + string(REGEX MATCH "^[0-9]+" _arch_num "${_arch}") + if(_arch_num AND _arch_num GREATER_EQUAL 90) + list(APPEND _mscclpp_ep_gpu_archs "${_arch}") + endif() endif() endforeach() - if(NOT PYBIND11_INCLUDE_DIR) + if(NOT _mscclpp_ep_gpu_archs) message(FATAL_ERROR - "pybind11 not found and not bundled with Torch. Install it (pip install pybind11) " - "or provide -Dpybind11_DIR=...") + "MSCCLPP_BUILD_EXT_EP=ON requires CUDA architecture 90 or newer. " + "Set MSCCLPP_GPU_ARCHS to include 90+ (for example, -DMSCCLPP_GPU_ARCHS=90).") endif() endif() -file(GLOB_RECURSE EP_SOURCES CONFIGURE_DEPENDS - buffer.cc +set(EP_SOURCES + moe_runtime.cc bindings.cpp - kernels/*.cu + kernels/low_latency.cu ) if(MSCCLPP_USE_ROCM) @@ -53,11 +44,11 @@ if(MSCCLPP_USE_ROCM) endif() # Build as a Python extension module (shared object with Python ABI suffix). -# The name `mscclpp_ep_cpp` matches the `TORCH_EXTENSION_NAME` hard-coded in -# `bindings.cpp` / `buffer.hpp`. -Python_add_library(mscclpp_ep_cpp MODULE ${EP_SOURCES}) +nanobind_add_module(mscclpp_ep_cpp ${EP_SOURCES}) +if(MSCCLPP_USE_CUDA) + set_target_properties(mscclpp_ep_cpp PROPERTIES CUDA_ARCHITECTURES "${_mscclpp_ep_gpu_archs}") +endif() -target_compile_definitions(mscclpp_ep_cpp PRIVATE TORCH_EXTENSION_NAME=mscclpp_ep_cpp) target_include_directories(mscclpp_ep_cpp PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${PROJECT_SOURCE_DIR}/include @@ -65,26 +56,6 @@ target_include_directories(mscclpp_ep_cpp PRIVATE ${PROJECT_SOURCE_DIR}/src/ext/include ${GPU_INCLUDE_DIRS} ) -if(pybind11_FOUND) - target_link_libraries(mscclpp_ep_cpp PRIVATE pybind11::module) -else() - target_include_directories(mscclpp_ep_cpp SYSTEM PRIVATE ${PYBIND11_INCLUDE_DIR}) -endif() - -target_link_libraries(mscclpp_ep_cpp PRIVATE ${TORCH_LIBRARIES}) -# Torch's CUDA interop library (ATen CUDAStream helpers used in buffer.cc). -if(TARGET torch::torch) - target_link_libraries(mscclpp_ep_cpp PRIVATE torch::torch) -endif() -# libtorch_python contains pybind11 bindings for torch::Tensor / torch::dtype -# (symbols like THPDtypeType). It is not listed in TORCH_LIBRARIES. -find_library(TORCH_PYTHON_LIBRARY torch_python - HINTS "${TORCH_INSTALL_PREFIX}/lib") -if(TORCH_PYTHON_LIBRARY) - target_link_libraries(mscclpp_ep_cpp PRIVATE ${TORCH_PYTHON_LIBRARY}) -else() - message(WARNING "libtorch_python not found; `import mscclpp_ep_cpp` may fail at runtime.") -endif() target_link_libraries(mscclpp_ep_cpp PRIVATE mscclpp ${GPU_LIBRARIES} Threads::Threads) # Number of intra-node NVLink peers compiled into the EP kernels. @@ -103,19 +74,6 @@ if(MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT) target_compile_definitions(mscclpp_ep_cpp PRIVATE MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT) endif() -# Experimental: NCCL-EP-style warp-specialized HT dispatch path. -# When ON, the internode `dispatch` kernel selects an alternate code path -# (guarded by `EP_DISPATCH_NCCLEP`) that ports NVIDIA NCCL-EP's -# warp-specialized overlap pipeline (concurrent inter-node fabric-write + -# intra-node drain) onto mscclpp MemoryChannel put/signal primitives, to -# close the GB200 MNNVL dispatch gap vs NCCL-EP. Default OFF = production path. -# -DMSCCLPP_EP_DISPATCH_NCCLEP=ON -option(MSCCLPP_EP_DISPATCH_NCCLEP - "Use the experimental NCCL-EP-ported warp-specialized HT dispatch kernel" OFF) -if(MSCCLPP_EP_DISPATCH_NCCLEP) - target_compile_definitions(mscclpp_ep_cpp PRIVATE EP_DISPATCH_NCCLEP) -endif() - set_target_properties(mscclpp_ep_cpp PROPERTIES PREFIX "" POSITION_INDEPENDENT_CODE ON diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 5d479cdf..b4ebd36c 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -1,45 +1,28 @@ # MSCCL++ Expert-Parallel (EP) extension -A port of DeepEP's MoE `dispatch` / `combine` primitives into MSCCL++, -targeting: - -- **High-Throughput (HT) mode** from [DeepEP](https://github.com/deepseek-ai/DeepEP), - branch `chhwang/dev-atomic-add-cleanup` — which already swaps NVSHMEM for - `mscclpp::PortChannel` / `mscclpp::MemoryChannel`. -- **Low-Latency (LL) mode** from [`nccl/contrib/nccl_ep`](https://github.com/NVIDIA/nccl/tree/master/contrib/nccl_ep), - which implements pure-RDMA dispatch/combine on top of the NCCL Device API. +A port of low-latency MoE `dispatch` / `combine` primitives into MSCCL++. +The active build exposes the LL `MoECommunicator` path. The older +High-Throughput (HT) prototype has been moved under `src/ext/ep/ht/` for +reference and is not compiled into `mscclpp_ep_cpp`. ## Status | Feature | Status | |------------------------------------|---------------------------------------------| -| `Buffer` construction + IPC + sync | ✅ ported (NVLink + RDMA) | -| `get_dispatch_layout` | ✅ ported | -| `intranode_dispatch` (NVLink) | ✅ validated (8 ranks H100; 4 ranks GB200) | -| `intranode_combine` (NVLink) | ✅ validated (8 ranks H100; 4 ranks GB200) | -| `internode_dispatch` (NVLink+RDMA) | ✅ validated (16 ranks 2×H100×8; 64 ranks 16×GB200) | -| `internode_combine` (NVLink+RDMA) | ✅ validated (16 ranks 2×H100×8; 64 ranks 16×GB200) | +| Runtime construction + IPC + sync | ✅ ported for LL | +| HT prototype | Archived under `src/ext/ep/ht/`; not compiled | | `low_latency_dispatch` (RDMA+IPC) | ✅ validated (8 ranks H100; 16 ranks 2×H100×8; 64 ranks 16×GB200 via NVLS fabric IPC) | | `low_latency_combine` (RDMA+IPC) | ✅ validated (8 ranks H100; 16 ranks 2×H100×8; 64 ranks 16×GB200 via NVLS fabric IPC) | | GB200 NVLS multimem fast path | ✅ runtime-gated by `mscclpp::isNvlsSupported()` | | Multi-`ProxyService` sharding | ✅ env-tunable, arch-aware default | | `Connection::atomicAdd` API | ✅ cherry-picked into mscclpp | -| Python frontend `mscclpp.ext.ep` | ✅ wraps HT + LL paths | -| pybind11 module `mscclpp_ep_cpp` | ✅ builds conditionally | +| Python frontend `mscclpp.ext.ep` | ✅ wraps the new LL MoECommunicator path | +| nanobind module `mscclpp_ep_cpp` | ✅ builds conditionally | -Internode HT was validated end-to-end on two H100×8 nodes connected over -Infiniband using [`test/python/ext/ep/test_internode_multirank.py`](../../../test/python/ext/ep/test_internode_multirank.py). -All 16 ranks complete dispatch followed by combine with exact (zero-diff) -recovery of the per-rank token payloads. - -On Azure GB200 NVL72 (4 GPUs / NUMA host, CX-7 RoCE), HT and LL were -validated at 16 nodes × 4 GPUs = **64 ranks** with HIDDEN=7168, -tokens=4096, experts=256, top-k=8: - -- HT internode: dispatch ~**2 006 GB/s** agg, combine ~**2 011 GB/s** agg - (`NVL_SEND=8 NVL_RECV=256 RDMA_SEND=8 RDMA_RECV=32`). -- LL internode: dispatch ~**16 817 GB/s** agg (~262 GB/s per rank), - combine ~**21 148 GB/s** agg (defaults). +On Azure GB200 NVL72 (4 GPUs / NUMA host, CX-7 RoCE), LL was validated at +16 nodes × 4 GPUs = **64 ranks** with HIDDEN=7168, tokens=4096, +experts=256, top-k=8: dispatch ~**16 817 GB/s** agg (~262 GB/s per rank), +combine ~**21 148 GB/s** agg (defaults). The GB200 path bypasses Azure CX-7 RoCE's broken `IBV_ATOMIC_*` by routing peer pointers through cuMem fabric IPC over the NVL72 fabric @@ -76,7 +59,7 @@ LL was validated on: 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 LL traffic. `Buffer` therefore allocates `N` +bottleneck for cross-node LL traffic. The EP runtime therefore allocates `N` ProxyServices and shards `PortChannel`s across them by `(qp_idx, dst_rank)`. @@ -92,41 +75,41 @@ dst_rank)`. peers go through a CPU proxy. The port is for functional parity, not latency. (Intra-node LL traffic uses CUDA IPC and is competitive.) - Unlike DeepEP, this port drives LL through `PortChannel` / - `MemoryChannel` rather than NVSHMEM, so `Buffer::sync()` connects + `MemoryChannel` rather than NVSHMEM, so runtime sync connects every peer even in `low_latency_mode=True`. -- The internode HT functional test inserts an explicit - `torch.cuda.synchronize()` + `dist.barrier()` between dispatch and - combine. Without it, fast ranks can launch combine while peers still - have in-flight dispatch proxy traffic, deadlocking the combine NVL - forwarder. Folding this barrier into - `Buffer::internode_dispatch` / `Buffer::internode_combine` (or - `cached_notify`) is tracked in the test's `XXX` comment. +- HT code is archived for reference only and is not part of the active + extension build. ## Build -The extension is **off by default** and requires PyTorch's CMake package: +For Python installs, use the `ep` extra: + +```bash +python -m pip install ".[cuda12,ep]" +``` + +The EP extension targets CUDA architectures **90 or newer**. Plain CMake builds +can enable it explicitly: ```bash -TORCH_CMAKE=$(python -c 'import torch; print(torch.utils.cmake_prefix_path)') cmake -S . -B build \ - -DMSCCLPP_BUILD_EXT_EP=ON \ - -DCMAKE_PREFIX_PATH="${TORCH_CMAKE}" + -DMSCCLPP_BUILD_EXT_EP=ON cmake --build build -j ``` -This produces `mscclpp_ep_cpp.so` — a pybind11 PyTorch extension module. +This produces `mscclpp_ep_cpp.so` — a nanobind extension module. The Python frontend picks it up automatically: ```python from mscclpp.ext import ep -buf = ep.Buffer(group, num_nvl_bytes=..., num_rdma_bytes=...) +moe_comm = ep.MoECommunicator(...) ``` ### Build-time CMake options | Variable | Default | Meaning | |---------------------------------------|---------|---------------------------------------------------------------| -| `MSCCLPP_BUILD_EXT_EP` | `OFF` | Build the EP extension at all | +| `MSCCLPP_BUILD_EXT_EP` | `ON` in Python wheels | Build the EP extension at all | | `MSCCLPP_EP_NUM_MAX_NVL_PEERS` | `8` | Compile-time `NUM_MAX_NVL_PEERS` — set to `4` for GB200 NVL72 | | `MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT` | `OFF` | Use a short ~10s kernel spin timeout (default is ~100s) | @@ -142,14 +125,11 @@ build-time settings are required: # Option 1: plain CMake. cmake -S . -B build \ -DMSCCLPP_BUILD_EXT_EP=ON \ - -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=4 \ - -DCMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')" + -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=4 cmake --build build -j -# Option 2: wheel-based install (pyproject.toml already sets -# MSCCLPP_BUILD_EXT_EP=ON). -CMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')" \ -python3 -m pip install --no-build-isolation \ +# Option 2: wheel-based install. +python3 -m pip install ".[cuda12,ep]" \ --config-settings=cmake.define.MSCCLPP_EP_NUM_MAX_NVL_PEERS=4 \ . ``` @@ -166,7 +146,7 @@ Runtime prerequisites on GB200: (`POSIX_FD | FABRIC`) can be exchanged across nodes. - NVLink-SHARP / multicast support enabled (`nvidia-smi mig … --imex` reachable). `mscclpp::isNvlsSupported()` must return `true` at - Buffer construction; otherwise the kernels fall back to the legacy + runtime construction; otherwise the kernels fall back to the legacy PortChannel + RDMA path (and on Azure CX-7 RoCE the broken IB atomics will hang). - `MSCCLPP_EP_LOCAL_WORLD_SIZE` partitions ranks into NUMA hosts; it @@ -268,33 +248,38 @@ topk=8, experts=256): > `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.Buffer` construction. +> built-in IB probe does not crash `ep.MoECommunicator` construction. ## Layout ``` src/ext/ep/ -├── CMakeLists.txt — builds mscclpp_ep_cpp (Torch + pybind11) +├── CMakeLists.txt — builds mscclpp_ep_cpp (nanobind) ├── README.md — this file -├── buffer.hpp / buffer.cc — host-side Buffer, sync(), dispatch/combine -├── config.hpp / event.hpp — Config, EventHandle -├── bindings.cpp — PYBIND11_MODULE definition +├── 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/ — archived HT prototype; not compiled +│ ├── buffer.hpp / buffer.cc +│ ├── config.hpp +│ ├── event.hpp +│ └── kernels/ +│ ├── buffer.cuh +│ ├── runtime.cu +│ ├── intranode_kernel.cu +│ ├── internode.cu +│ └── internode_ncclep.cuh └── kernels/ ├── api.cuh — host-callable kernel prototypes ├── configs.cuh — compile-time constants (GPU-only) - ├── buffer.cuh — Buffer/AsymBuffer/SymBuffer helpers ├── exception.cuh — EP_HOST/DEVICE_ASSERT + CUDA_CHECK ├── launch.cuh — SETUP_LAUNCH_CONFIG / SWITCH_* macros ├── utils.cuh — device inline helpers - ├── runtime.cu — intranode::barrier launcher - ├── intranode_kernel.cu — intranode dispatch/combine kernels - ├── internode.cu — internode HT dispatch/combine + layout - │ (incl. NVLS multimem fast path for GB200) - └── internode_ll.cu — internode LL dispatch/combine + └── low_latency.cu — LL dispatch/combine (RDMA + IPC paths) python/mscclpp/ext/ep/ -├── __init__.py — reexports Buffer / Config / EventHandle -└── buffer.py — torch.distributed-aware frontend +├── __init__.py — reexports the public MoECommunicator API +└── communicator.py — torch.Tensor frontend over raw-pointer runtime calls test/python/ext/ep/ ├── test_intranode_multirank.py — intranode HT dispatch+combine @@ -321,7 +306,7 @@ 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 pybind11 blake3 sortedcontainers +conda install -c conda-forge -y cupy mpi4py nanobind blake3 sortedcontainers # PyTorch (pulls a matching cuda-toolkit + NCCL). pip3 install torch @@ -445,25 +430,29 @@ Env knobs: | `MSCCLPP_EP_BENCH_HIDDEN` | Hidden dim | `7168` | | `MSCCLPP_EP_BENCH_EXPERTS` | Total experts | test-specific | | `MSCCLPP_EP_BENCH_TOPK` | top-k routing | `8` | -| `MSCCLPP_EP_NUM_PROXIES` | Number of `ProxyService`s in `Buffer` | 8 (Hopper) / 1 (Blackwell) | +| `MSCCLPP_EP_NUM_PROXIES` | Number of runtime `ProxyService`s | 8 (Hopper) / 1 (Blackwell) | ## Migration plan +This section is historical. The HT prototype from the earlier migration phases +is archived under `src/ext/ep/ht/` and is not part of the active extension +target; the current `mscclpp_ep_cpp` build compiles the LL runtime only. + ### Phase 1 — DONE - [x] Copy DeepEP kernel headers (configs / buffer / utils / launch / exception). - [x] Port intranode kernels + runtime (NVLink only). - [x] Port `get_dispatch_layout` (host-safe subset of internode kernels). -- [x] Port host Buffer: ctor, sync, get_dispatch_layout, intranode +- [x] Port host runtime: ctor, sync, get_dispatch_layout, intranode dispatch/combine. -- [x] pybind11 `mscclpp_ep_cpp` module + Python frontend. +- [x] nanobind `mscclpp_ep_cpp` module + Python frontend. ### Phase 2 — internode HT (NVLink + RDMA) — DONE - [x] Port `notify_dispatch`, `dispatch`, `cached_notify`, `combine` kernels. -- [x] Wire `Buffer::internode_dispatch` / `Buffer::internode_combine` host +- [x] Wire internode dispatch/combine host orchestration. -- [x] `Buffer::sync()` builds `port_channel_handles_device_ptr` and +- [x] Runtime sync builds `port_channel_handles_device_ptr` and `memory_channel_handles_device_ptr`, with port channels ordered by peer rank (so kernel-side indexing by `peer_rank` is consistent). - [x] Validated on 2×H100×8 with `test_internode_multirank.py`. @@ -472,7 +461,7 @@ Env knobs: ### Phase 3 — Low-Latency (RDMA + CUDA-IPC) — DONE -Port `DeepEP/csrc/kernels/internode_ll.cu` and cross-reference +Port `DeepEP/csrc/kernels/internode_ll.cu` as `kernels/low_latency.cu` and cross-reference `nccl/contrib/nccl_ep/device/low_latency.cu`. The nccl_ep reference is modular (see `device_primitives.cuh`, `hybrid_ep.cuh`) and uses NCCL Device API; the translation table is: @@ -488,7 +477,7 @@ Device API; the translation table is: | NVSHMEM symmetric heap | `cudaMalloc` + `ProxyService::addMemory` | | NVSHMEM barrier | `bootstrap->barrier()` or `intranode::barrier` | -`Buffer::low_latency_dispatch` / `low_latency_combine` are validated +Low-latency dispatch/combine are validated intra-node (8 ranks, CUDA IPC fast path) and cross-node (16 ranks on 2×H100×8, IPC + IB). Functional correctness is bit-exact against the reference dispatch/combine. @@ -514,7 +503,7 @@ on top of the H100 baseline: version-guarded, `NUM_TIMEOUT_CYCLES` restored, CMake install destination dual-mode. - [x] **LL bypass for broken CX-7 IB atomics** (Proposal A): - `Buffer::rdma_buffer_ptr` allocated via + runtime RDMA buffer allocated via `mscclpp::detail::gpuCallocPhysical` (POSIX_FD | FABRIC handles); LL IPC fast-path gate lifted from `num_rdma_ranks==1` to `low_latency_mode`; peer bases resolved through diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index bf3143a5..2d6e84f8 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -4,79 +4,90 @@ // Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP) // branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. // -// pybind11 module definition for the MSCCL++ EP extension. Mirrors -// DeepEP's `PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)` so call sites port -// with minimal changes. +// nanobind module definition for the MSCCL++ EP extension. -#include -#include -#include -#include +#include +#include +#include + +#include +#include +#include +#include -#include "buffer.hpp" #include "config.hpp" +#include "kernels/api.cuh" +#include "moe_runtime.hpp" -namespace py = pybind11; +namespace nb = nanobind; -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +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); } + +} // namespace + +NB_MODULE(mscclpp_ep_cpp, m) { m.doc() = "MSCCL++ Expert-Parallel (MoE dispatch/combine) extension"; - py::class_(m, "Config") - .def(py::init(), py::arg("num_sms") = 20, py::arg("num_max_nvl_chunked_send_tokens") = 6, - py::arg("num_max_nvl_chunked_recv_tokens") = 256, py::arg("num_max_rdma_chunked_send_tokens") = 6, - py::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); + m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::getLowLatencyRdmaSizeHint); - m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::get_low_latency_rdma_size_hint); + nb::module_::import_("mscclpp._mscclpp"); - py::class_(m, "EventHandle") - .def(py::init<>()) - .def("current_stream_wait", &mscclpp::ep::EventHandle::current_stream_wait); + nb::enum_(m, "MoEMode") + .value("LOW_LATENCY", mscclpp::ep::MoEMode::LOW_LATENCY) + .value("HIGH_THROUGHPUT", mscclpp::ep::MoEMode::HIGH_THROUGHPUT); - // NOTE: `mscclpp::UniqueId` is `std::array`, which pybind11 - // implicitly converts to a Python list. We therefore avoid exposing it as - // a py::class_ and convert to/from `py::bytes` at the binding boundary. + nb::enum_(m, "DispatchLayout") + .value("EXPERT_MAJOR", mscclpp::ep::DispatchLayout::EXPERT_MAJOR) + .value("FLAT", mscclpp::ep::DispatchLayout::FLAT); - py::class_(m, "Buffer") - .def(py::init(), py::arg("rank"), py::arg("num_ranks"), - py::arg("num_nvl_bytes"), py::arg("num_rdma_bytes"), py::arg("low_latency_mode")) - .def("is_available", &mscclpp::ep::Buffer::is_available) - .def("is_internode_available", &mscclpp::ep::Buffer::is_internode_available) - .def("get_num_rdma_ranks", &mscclpp::ep::Buffer::get_num_rdma_ranks) - .def("get_rdma_rank", &mscclpp::ep::Buffer::get_rdma_rank) - .def("get_root_rdma_rank", &mscclpp::ep::Buffer::get_root_rdma_rank) - .def("get_local_device_id", &mscclpp::ep::Buffer::get_local_device_id) - .def("get_local_ipc_handle", &mscclpp::ep::Buffer::get_local_ipc_handle) - .def("get_local_nvshmem_unique_id", &mscclpp::ep::Buffer::get_local_nvshmem_unique_id) - .def("get_local_buffer_tensor", &mscclpp::ep::Buffer::get_local_buffer_tensor) - .def("create_unique_id", - [](const mscclpp::ep::Buffer& self) { - auto uid = self.create_unique_id(); - return py::bytes(reinterpret_cast(uid.data()), uid.size()); - }) - .def("connect", - [](mscclpp::ep::Buffer& self, py::bytes data) { - std::string s = data; - mscclpp::UniqueId uid; - if (s.size() != uid.size()) { - throw std::runtime_error("mscclpp_ep_cpp.Buffer.connect: UniqueId size mismatch"); - } - std::memcpy(uid.data(), s.data(), s.size()); - self.connect(uid); - }) - .def("sync", &mscclpp::ep::Buffer::sync) - .def("get_dispatch_layout", &mscclpp::ep::Buffer::get_dispatch_layout) - .def("intranode_dispatch", &mscclpp::ep::Buffer::intranode_dispatch) - .def("intranode_combine", &mscclpp::ep::Buffer::intranode_combine) - .def("internode_dispatch", &mscclpp::ep::Buffer::internode_dispatch) - .def("internode_combine", &mscclpp::ep::Buffer::internode_combine) - .def("clean_low_latency_buffer", &mscclpp::ep::Buffer::clean_low_latency_buffer) - .def("low_latency_dispatch", &mscclpp::ep::Buffer::low_latency_dispatch, py::arg("x"), py::arg("topk_idx"), - py::arg("num_max_dispatch_tokens_per_rank"), py::arg("num_experts"), py::arg("use_fp8"), py::arg("async"), - py::arg("return_recv_hook"), py::arg("out_packed_recv_x") = py::none(), - py::arg("out_packed_recv_x_scales") = py::none(), py::arg("out_packed_recv_src_info") = py::none(), - py::arg("out_packed_recv_layout_range") = py::none(), py::arg("out_packed_recv_count") = py::none()) - .def("low_latency_combine", &mscclpp::ep::Buffer::low_latency_combine) - .def("get_next_low_latency_combine_buffer", &mscclpp::ep::Buffer::get_next_low_latency_combine_buffer); + nb::class_(m, "MoERuntime") + .def(nb::init(), nb::arg("comm"), + nb::arg("num_nvl_bytes"), nb::arg("num_rdma_bytes"), nb::arg("mode")) + .def("is_available", &mscclpp::ep::MoERuntime::isAvailable) + .def("is_internode_available", &mscclpp::ep::MoERuntime::isInternodeAvailable) + .def("get_num_rdma_ranks", &mscclpp::ep::MoERuntime::getNumRdmaRanks) + .def("get_rdma_rank", &mscclpp::ep::MoERuntime::getRdmaRank) + .def("get_root_rdma_rank", &mscclpp::ep::MoERuntime::getRootRdmaRank) + .def("get_local_device_id", &mscclpp::ep::MoERuntime::getLocalDeviceId) + .def("get_local_ipc_handle", + [](const mscclpp::ep::MoERuntime& self) { return stringToBytes(self.getLocalIpcHandle()); }) + .def( + "dispatch", + [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t outputPtr, + uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, + uintptr_t outputCountPtr, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, + int numExperts, bool requiresQuantization, mscclpp::ep::DispatchLayout outputLayout, uintptr_t streamPtr) { + self.dispatch( + ptr(outputPtr), reinterpret_cast(ptr(outputScalesPtr)), + reinterpret_cast(ptr(outputSrcInfoPtr)), reinterpret_cast(ptr(outputLayoutRangePtr)), + reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), reinterpret_cast(ptr(topkIdxPtr)), + numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, numExperts, requiresQuantization, outputLayout, + stream(streamPtr)); + }, + nb::arg("input_ptr"), nb::arg("topk_idx_ptr"), nb::arg("output_ptr"), nb::arg("output_scales_ptr"), + nb::arg("output_src_info_ptr"), nb::arg("output_layout_range_ptr"), nb::arg("output_count_ptr"), + nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("num_max_dispatch_tokens_per_rank"), + nb::arg("num_experts"), nb::arg("requires_quantization"), nb::arg("output_layout"), nb::arg("stream_ptr")) + .def( + "combine", + [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t expertScalesPtr, uintptr_t topkIdxPtr, + uintptr_t topkWeightsPtr, uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, + int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, + bool requiresDequantization, uintptr_t streamPtr) { + self.combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(expertScalesPtr)), + reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), + reinterpret_cast(ptr(srcInfoPtr)), reinterpret_cast(ptr(layoutRangePtr)), + numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, numExperts, requiresDequantization, + stream(streamPtr)); + }, + nb::arg("expert_output_ptr"), nb::arg("expert_scales_ptr"), nb::arg("topk_idx_ptr"), + nb::arg("topk_weights_ptr"), nb::arg("src_info_ptr"), nb::arg("layout_range_ptr"), nb::arg("output_ptr"), + nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("num_max_dispatch_tokens_per_rank"), + nb::arg("num_experts"), nb::arg("requires_dequantization"), nb::arg("stream_ptr")); } diff --git a/src/ext/ep/config.hpp b/src/ext/ep/config.hpp index 5bff86d5..02897b4f 100644 --- a/src/ext/ep/config.hpp +++ b/src/ext/ep/config.hpp @@ -2,157 +2,52 @@ // Licensed under the MIT License. #pragma once -#include "kernels/api.cuh" +#include +#include +#include + +#include "kernels/configs.cuh" #include "kernels/exception.cuh" namespace mscclpp { namespace ep { template -dtype_t cell_div(dtype_t a, dtype_t b) { +dtype_t cellDiv(dtype_t a, dtype_t b) { return (a + b - 1) / b; } template dtype_t align(dtype_t a, dtype_t b) { - return cell_div(a, b) * b; + return cellDiv(a, b) * b; } -struct 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, - int num_max_rdma_chunked_send_tokens, int num_max_rdma_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); - 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 = - align(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 = num_sms / 2; - - 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 - // Increment 3 (cross-GPU peer-map): a fixed-capacity recv-output pool appended - // to the NVL cudaMalloc (AFTER the fifo/task regions, NOT added to - // num_nvl_bytes, so it does not count against the INT_MAX-limited registered - // size). Peers reach it through the same IPC handle as buffer_ptrs. Holds a - // small header (the local recv_gbl_rank_prefix_sum, made peer-readable so a - // cross-GPU forwarder can compute the destination's final recv_x index) - // followed by the recv_x hidden region. recv_x becomes a zero-copy view of - // that region; the cross-GPU forwarder writes hidden straight to the - // destination peer's pool, removing the receiver's hidden drain. Falls back to - // torch::empty + receiver when num_recv_tokens exceeds the capacity. - static constexpr int kEpRecvPoolMaxTokens = 65536; - static constexpr int64_t kEpRecvPoolMaxHiddenBytes = 16384; // worst-case per-token bytes (hidden<=8192 bf16) - size_t get_recv_pool_header_bytes(int num_ranks) const { - return ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - } - // Total bytes of the recv-output pool. Sized for the worst-case hidden so the - // Buffer ctor can allocate it without knowing the runtime hidden dim. - static size_t recv_pool_bytes_static(int num_ranks) { - size_t header = ((static_cast(num_ranks) * sizeof(int) + 127) / 128) * 128; - size_t b = header + static_cast(kEpRecvPoolMaxTokens) * static_cast(kEpRecvPoolMaxHiddenBytes); - 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); - } - - 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; - - // 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 = num_sms / 2; - - 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; - // Two extra uint64_t scratch slots per (channel, rdma_rank) used by the - // dispatch/combine kernels as the RDMA WRITE source for absolute-value - // tail/head updates (replaces broken HW atomicAdd on Azure CX-7 RoCE). - num_bytes += num_channels * num_rdma_ranks * sizeof(uint64_t) * 2; - num_bytes = ((num_bytes + 127) / 128) * 128; - return num_bytes; - } -}; - struct LowLatencyBuffer { - int num_clean_int = 0; + int numCleanInt = 0; - void* dispatch_rdma_send_buffer = nullptr; - void* dispatch_rdma_recv_data_buffer = nullptr; + void* dispatchRdmaSendBuffer = nullptr; + void* dispatchRdmaRecvDataBuffer = nullptr; // NOTE: signaling buffers are int64_t (not int) so that IB atomic ops // (IBV_WR_ATOMIC_FETCH_AND_ADD is a 64-bit, 8-byte-aligned op) always // target an 8-byte-aligned address. Using int32 slots produced unaligned // atomics at odd indices that the NIC silently drops. - int64_t* dispatch_rdma_recv_count_buffer = nullptr; + int64_t* dispatchRdmaRecvCountBuffer = nullptr; - void* combine_rdma_send_buffer = nullptr; - void* combine_rdma_recv_data_buffer = nullptr; - int64_t* combine_rdma_recv_flag_buffer = nullptr; + void* combineRdmaSendBuffer = nullptr; + void* combineRdmaRecvDataBuffer = nullptr; + int64_t* combineRdmaRecvFlagBuffer = nullptr; - void* combine_rdma_send_buffer_data_start = nullptr; - size_t num_bytes_per_combine_msg = 0; + void* combineRdmaSendBufferDataStart = nullptr; + size_t numBytesPerCombineMsg = 0; - std::pair clean_meta() { - EP_HOST_ASSERT(dispatch_rdma_recv_count_buffer == combine_rdma_recv_flag_buffer); - return {dispatch_rdma_recv_count_buffer, num_clean_int}; + std::pair cleanMeta() { + EP_HOST_ASSERT(dispatchRdmaRecvCountBuffer == combineRdmaRecvFlagBuffer); + return {dispatchRdmaRecvCountBuffer, numCleanInt}; } }; struct LowLatencyLayout { - size_t total_bytes = 0; + size_t totalBytes = 0; LowLatencyBuffer buffers[2]; template @@ -160,9 +55,9 @@ struct LowLatencyLayout { return reinterpret_cast(reinterpret_cast(ptr) + count); } - LowLatencyLayout(void* rdma_buffer, int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, - int num_experts) { - const int num_scales = hidden / 128; + LowLatencyLayout(void* rdmaBuffer, int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts) { + (void)numRanks; + const int numScales = hidden / 128; // Dispatch and combine layout: // - 2 symmetric odd/even send buffer @@ -171,56 +66,52 @@ struct LowLatencyLayout { // Message sizes // NOTES: you should add a control `int4` for combine messages if you want to do data transformation - EP_HOST_ASSERT(num_scales * sizeof(float) <= hidden); - size_t num_bytes_per_dispatch_msg = - sizeof(int4) + std::max(hidden * sizeof(nv_bfloat16), hidden + num_scales * sizeof(float)); - size_t num_bytes_per_combine_msg = hidden * sizeof(nv_bfloat16); + EP_HOST_ASSERT(numScales * static_cast(sizeof(float)) <= hidden); + size_t numBytesPerDispatchMsg = + sizeof(int4) + std::max(hidden * sizeof(nv_bfloat16), hidden + numScales * sizeof(float)); + size_t numBytesPerCombineMsg = hidden * sizeof(nv_bfloat16); // Send buffer - size_t dispatch_send_buffer_bytes = num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg; - size_t combine_send_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg; - size_t send_buffer_bytes = std::max(dispatch_send_buffer_bytes, combine_send_buffer_bytes); - EP_HOST_ASSERT(send_buffer_bytes % sizeof(int4) == 0); - total_bytes += send_buffer_bytes * 2; + size_t dispatchSendBufferBytes = numMaxDispatchTokensPerRank * numBytesPerDispatchMsg; + size_t combineSendBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; + size_t sendBufferBytes = std::max(dispatchSendBufferBytes, combineSendBufferBytes); + EP_HOST_ASSERT(sendBufferBytes % sizeof(int4) == 0); + totalBytes += sendBufferBytes * 2; // Symmetric receive buffers // TODO: optimize memory usages - size_t dispatch_recv_data_buffer_bytes = - num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg; - size_t combine_recv_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg; - size_t recv_buffer_bytes = std::max(dispatch_recv_data_buffer_bytes, combine_recv_buffer_bytes); - EP_HOST_ASSERT(recv_buffer_bytes % sizeof(int4) == 0); - total_bytes += recv_buffer_bytes * 2; + size_t dispatchRecvDataBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerDispatchMsg; + size_t combineRecvBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; + size_t recvBufferBytes = std::max(dispatchRecvDataBufferBytes, combineRecvBufferBytes); + EP_HOST_ASSERT(recvBufferBytes % sizeof(int4) == 0); + totalBytes += recvBufferBytes * 2; // Symmetric signaling buffers (int64_t slots for 8-byte-aligned IB atomics). - size_t dispatch_recv_count_buffer_bytes = num_experts * sizeof(int64_t); - size_t combine_recv_flag_buffer_bytes = dispatch_recv_count_buffer_bytes; - size_t signaling_buffer_bytes = std::max(dispatch_recv_count_buffer_bytes, combine_recv_flag_buffer_bytes); - total_bytes += signaling_buffer_bytes * 2; + size_t dispatchRecvCountBufferBytes = numExperts * sizeof(int64_t); + size_t combineRecvFlagBufferBytes = dispatchRecvCountBufferBytes; + size_t signalingBufferBytes = std::max(dispatchRecvCountBufferBytes, combineRecvFlagBufferBytes); + totalBytes += signalingBufferBytes * 2; // Assign pointers // NOTES: we still leave some space for distinguishing dispatch/combine buffer, // so you may see some parameters are duplicated for (int i = 0; i < 2; ++i) { - buffers[i] = { - static_cast(signaling_buffer_bytes / sizeof(int64_t)), - advance(rdma_buffer, send_buffer_bytes * i), - advance(rdma_buffer, send_buffer_bytes * 2 + recv_buffer_bytes * i), - advance(rdma_buffer, send_buffer_bytes * 2 + recv_buffer_bytes * 2 + signaling_buffer_bytes * i), - advance(rdma_buffer, send_buffer_bytes * i), - advance(rdma_buffer, send_buffer_bytes * 2 + recv_buffer_bytes * i), - advance(rdma_buffer, send_buffer_bytes * 2 + recv_buffer_bytes * 2 + signaling_buffer_bytes * i), - advance(rdma_buffer, send_buffer_bytes * i), - num_bytes_per_combine_msg}; + buffers[i] = {static_cast(signalingBufferBytes / sizeof(int64_t)), + advance(rdmaBuffer, sendBufferBytes * i), + advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * i), + advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * 2 + signalingBufferBytes * i), + advance(rdmaBuffer, sendBufferBytes * i), + advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * i), + advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * 2 + signalingBufferBytes * i), + advance(rdmaBuffer, sendBufferBytes * i), + numBytesPerCombineMsg}; } } }; -inline size_t get_low_latency_rdma_size_hint(int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, - int num_experts) { - auto num_bytes = - LowLatencyLayout(nullptr, num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts).total_bytes; - return ((num_bytes + NUM_BUFFER_ALIGNMENT_BYTES) / NUM_BUFFER_ALIGNMENT_BYTES) * NUM_BUFFER_ALIGNMENT_BYTES; +inline size_t getLowLatencyRdmaSizeHint(int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts) { + auto numBytes = LowLatencyLayout(nullptr, numMaxDispatchTokensPerRank, hidden, numRanks, numExperts).totalBytes; + return ((numBytes + NUM_BUFFER_ALIGNMENT_BYTES) / NUM_BUFFER_ALIGNMENT_BYTES) * NUM_BUFFER_ALIGNMENT_BYTES; } } // namespace ep diff --git a/src/ext/ep/buffer.cc b/src/ext/ep/ht/buffer.cc similarity index 90% rename from src/ext/ep/buffer.cc rename to src/ext/ep/ht/buffer.cc index e76bed9c..98dda128 100644 --- a/src/ext/ep/buffer.cc +++ b/src/ext/ep/ht/buffer.cc @@ -3,8 +3,6 @@ #include #include #include -#include -#include #include #include @@ -12,8 +10,8 @@ #include #include -#include "kernels/api.cuh" -#include "kernels/configs.cuh" +#include "../kernels/api.cuh" +#include "../kernels/configs.cuh" namespace mscclpp { namespace ep { @@ -308,11 +306,9 @@ int Buffer::get_root_rdma_rank(bool global) const { return global ? nvl_rank : 0 int Buffer::get_local_device_id() const { return device_id; } -pybind11::bytearray Buffer::get_local_ipc_handle() const { - return {ipc_handles[nvl_rank].reserved, CUDA_IPC_HANDLE_SIZE}; -} +std::string Buffer::get_local_ipc_handle() const { return {ipc_handles[nvl_rank].reserved, CUDA_IPC_HANDLE_SIZE}; } -pybind11::bytearray Buffer::get_local_nvshmem_unique_id() const { +std::string Buffer::get_local_nvshmem_unique_id() 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. @@ -320,14 +316,11 @@ pybind11::bytearray Buffer::get_local_nvshmem_unique_id() const { "mscclpp::ep::Buffer::get_local_nvshmem_unique_id: not applicable (NVSHMEM is not used in mscclpp_ep)"); } -torch::Tensor Buffer::get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, - bool use_rdma_buffer) const { - torch::ScalarType casted_dtype = torch::python::detail::py_object_to_dtype(dtype); - auto element_bytes = static_cast(elementSize(casted_dtype)); +torch::Tensor Buffer::get_local_buffer_tensor(torch::ScalarType dtype, int64_t offset, bool use_rdma_buffer) const { + auto element_bytes = static_cast(elementSize(dtype)); auto base_ptr = reinterpret_cast(use_rdma_buffer ? rdma_buffer_ptr : buffer_ptrs[nvl_rank]) + offset; auto num_bytes = use_rdma_buffer ? num_rdma_bytes : num_nvl_bytes; - return torch::from_blob(base_ptr, num_bytes / element_bytes, - torch::TensorOptions().dtype(casted_dtype).device(at::kCUDA)); + return torch::from_blob(base_ptr, num_bytes / element_bytes, torch::TensorOptions().dtype(dtype).device(at::kCUDA)); } mscclpp::UniqueId Buffer::create_unique_id() const { return bootstrap->createUniqueId(); } @@ -338,8 +331,8 @@ void Buffer::connect(mscclpp::UniqueId root_id) { } void Buffer::sync(const std::vector& device_ids, - const std::vector>& all_gathered_handles, - const std::optional& root_unique_id_opt) { + const std::vector>& all_gathered_handles, + const std::optional& root_unique_id_opt) { EP_HOST_ASSERT(not is_available()); const std::vector ib_transports = { @@ -355,7 +348,7 @@ void Buffer::sync(const std::vector& device_ids, EP_HOST_ASSERT(device_ids.size() == all_gathered_handles.size()); for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++i) { EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value()); - auto handle_str = std::string(all_gathered_handles[offset + i].value()); + const auto& handle_str = all_gathered_handles[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); @@ -744,15 +737,16 @@ void Buffer::sync(const std::vector& device_ids, // `nodes_config.cfg` covering all node IPs. // ------------------------------------------------------------------ // ------------------------------------------------------------------ - // Phase 4: HT internode mode also benefits from fabric-IPC peer - // pointers. On Azure CX-7 the IB RDMA WRITE that PortChannel uses - // for the dispatch/combine token data payload hangs cross-node - // (same root cause as signal/wait), so we set up the same per-peer - // mapped pointers as LL and the kernels write tokens directly into - // the peer's `rdma_buffer_ptr` via the NVL72 fabric VA. + // 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 bool want_fabric_ipc = use_fabric_ipc_alloc; - if (want_fabric_ipc) { + 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 @@ -763,7 +757,7 @@ void Buffer::sync(const std::vector& device_ids, 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) continue; + if (r == rank || !is_ipc_peer(r)) continue; communicator->sendMemory(rdma_mem_ipc, r, kLlIpcTag); remote_futures[r] = communicator->recvMemory(r, kLlIpcTag); } @@ -772,24 +766,24 @@ void Buffer::sync(const std::vector& device_ids, std::vector> conn_futures(num_ranks); mscclpp::EndpointConfig cfg(ipc_transport); for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; + 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) continue; + 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 - // and is dereferenceable from the GPU. + // 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) continue; + if (r == rank || !is_ipc_peer(r)) continue; remote_mems[r] = remote_futures[r].get(); peer_rdma_bases[r] = remote_mems[r].data(); } @@ -804,14 +798,15 @@ void Buffer::sync(const std::vector& device_ids, fflush(stdout); } - // 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. +#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) { @@ -836,8 +831,8 @@ void Buffer::sync(const std::vector& device_ids, 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)); + 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]); @@ -852,22 +847,24 @@ void Buffer::sync(const std::vector& device_ids, sizeof(int) * static_cast(Config::kEpRecvPoolMaxTokens) * num_ranks)); } } +#endif if (low_latency_mode) { // LL barrier ring needs MemoryChannels. - std::vector ll_handles(num_ranks); + std::vector ll_handles(num_ranks); for (int r = 0; r < num_ranks; ++r) { - if (r == rank) continue; + 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); + mscclpp::detail::gpuCallocShared(num_ranks); + mscclpp::gpuMemcpy( + ll_memory_channel_handles_device_ptr.get(), ll_handles.data(), num_ranks, cudaMemcpyHostToDevice); - ll_ipc_ready = true; + ll_ranks_per_ipc_domain = ipc_domain_size; + ll_ipc_ready = ipc_domain_size >= num_ranks; } } } @@ -1096,6 +1093,7 @@ Buffer::intranode_dispatch( // peer-mapped recv pool so the sender writes hidden straight to its final slot, // eliminating the 2-hop ring + receiver hidden drain. Falls back to torch::empty. 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 = @@ -1104,6 +1102,10 @@ Buffer::intranode_dispatch( static_cast(num_recv_tokens) * hidden * static_cast(x.element_size()) <= 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 // Allocate new tensors torch::Tensor recv_x; @@ -1321,7 +1323,6 @@ Buffer::internode_dispatch( const std::optional& cached_recv_gbl_rank_prefix_sum, int expert_alignment, const Config& config, std::optional& previous_event, bool async, bool allocate_on_comm_stream) { // In dispatch, CPU will busy-wait until GPU receive tensor size metadata from other ranks, which can be quite long. - pybind11::gil_scoped_release release; const int num_channels = config.num_sms / 2; EP_HOST_ASSERT(config.num_sms % 2 == 0); @@ -1503,7 +1504,7 @@ Buffer::internode_dispatch( } // Allocate new tensors - void** ep_recv_pool_ptrs = nullptr; // non-null selects the increment-4 VMM direct-write path + 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 @@ -1717,17 +1718,17 @@ std::tuple, std::optional(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( - at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), combined_x.data_ptr(), combined_topk_weights_ptr, - is_combined_token_in_rank.data_ptr(), x.data_ptr(), topk_weights_ptr, combined_rdma_head.data_ptr(), - combined_nvl_head.data_ptr(), src_meta.data_ptr(), rdma_channel_prefix_matrix.data_ptr(), - rdma_rank_prefix_sum.data_ptr(), gbl_channel_prefix_matrix.data_ptr(), num_tokens, num_combined_tokens, - 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, num_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); + internode::combine(at::cuda::ScalarTypeToCudaDataType(x.scalar_type()), combined_x.data_ptr(), + combined_topk_weights_ptr, is_combined_token_in_rank.data_ptr(), x.data_ptr(), + topk_weights_ptr, combined_rdma_head.data_ptr(), combined_nvl_head.data_ptr(), + src_meta.data_ptr(), rdma_channel_prefix_matrix.data_ptr(), + rdma_rank_prefix_sum.data_ptr(), gbl_channel_prefix_matrix.data_ptr(), num_tokens, + num_combined_tokens, 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, num_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); std::optional event; if (async) { @@ -1764,11 +1765,18 @@ void Buffer::clean_low_latency_buffer(int num_max_dispatch_tokens_per_rank, int check_boundary(clean_meta_0.first, clean_meta_0.second * sizeof(int)); check_boundary(clean_meta_1.first, clean_meta_1.second * sizeof(int)); - internode_ll::clean_low_latency_buffer( - clean_meta_0.first, clean_meta_0.second, clean_meta_1.first, clean_meta_1.second, rank, num_ranks, - port_channel_handles_device_ptr.get(), - ll_memory_channel_handles_device_ptr ? ll_memory_channel_handles_device_ptr.get() : nullptr, ll_ipc_ready, - at::cuda::getCurrentCUDAStream()); + low_latency::TransportContext transport{ + .rdmaBufferBase_ = rdma_buffer_ptr, + .portChannels_ = port_channel_handles_device_ptr.get(), + .memoryChannels_ = ll_memory_channel_handles_device_ptr ? ll_memory_channel_handles_device_ptr.get() : nullptr, + .peerBases_ = peer_rdma_bases_gpu, + .ipcReady_ = ll_ipc_ready, + .rank_ = rank, + .numRanks_ = num_ranks, + .ranksPerIpcDomain_ = ll_ranks_per_ipc_domain}; + + low_latency::cleanBuffers(clean_meta_0.first, clean_meta_0.second, clean_meta_1.first, clean_meta_1.second, transport, + at::cuda::getCurrentCUDAStream()); } std::tuple, torch::Tensor, torch::Tensor, torch::Tensor, @@ -1879,21 +1887,44 @@ Buffer::low_latency_dispatch(const torch::Tensor& x, const torch::Tensor& topk_i } auto next_clean_meta = next_buffer.clean_meta(); - auto port_handles = port_channel_handles_device_ptr.get(); - auto mem_handles = ll_memory_channel_handles_device_ptr ? ll_memory_channel_handles_device_ptr.get() : nullptr; - auto peer_bases = peer_rdma_bases_gpu; - const bool use_ipc = ll_ipc_ready; - auto rdma_base = rdma_buffer_ptr; - auto launcher = [=](int phases) { - internode_ll::dispatch(packed_recv_x.data_ptr(), packed_recv_x_scales_ptr, packed_recv_src_info.data_ptr(), - packed_recv_layout_range.data_ptr(), packed_recv_count.data_ptr(), - buffer.dispatch_rdma_recv_data_buffer, buffer.dispatch_rdma_recv_count_buffer, - buffer.dispatch_rdma_send_buffer, x.data_ptr(), topk_idx.data_ptr(), - next_clean_meta.first, next_clean_meta.second, num_tokens, hidden, - num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank, num_ranks, use_fp8, workspace, - launch_stream, phases, rdma_base, port_handles, peer_bases, mem_handles, use_ipc); + + low_latency::DispatchConfig config{.numTokens_ = num_tokens, + .hidden_ = hidden, + .numTopk_ = num_topk, + .numExperts_ = num_experts, + .numMaxTokensPerRank_ = num_max_dispatch_tokens_per_rank, + .inputDType_ = low_latency::DType::BF16, + .outputDType_ = use_fp8 ? low_latency::DType::F8E4M3 : low_latency::DType::BF16, + .outputLayout_ = low_latency::DispatchLayout::EXPERT_MAJOR}; + low_latency::BufferSet current_buf{.sendDataBuffer_ = buffer.dispatch_rdma_send_buffer, + .sendCountBuffer_ = nullptr, + .recvDataBuffer_ = buffer.dispatch_rdma_recv_data_buffer, + .recvCountBuffer_ = buffer.dispatch_rdma_recv_count_buffer, + .cleanupRegion_ = nullptr, + .cleanupSize_ = 0}; + low_latency::BufferSet next_buf{.sendDataBuffer_ = nullptr, + .sendCountBuffer_ = nullptr, + .recvDataBuffer_ = nullptr, + .recvCountBuffer_ = nullptr, + .cleanupRegion_ = next_clean_meta.first, + .cleanupSize_ = next_clean_meta.second}; + low_latency::TransportContext transport{ + .rdmaBufferBase_ = rdma_buffer_ptr, + .portChannels_ = port_channel_handles_device_ptr.get(), + .memoryChannels_ = ll_memory_channel_handles_device_ptr ? ll_memory_channel_handles_device_ptr.get() : nullptr, + .peerBases_ = peer_rdma_bases_gpu, + .ipcReady_ = ll_ipc_ready, + .rank_ = rank, + .numRanks_ = num_ranks, + .ranksPerIpcDomain_ = ll_ranks_per_ipc_domain}; + + auto launcher = [=](low_latency::Phase phase) { + low_latency::dispatch(packed_recv_x.data_ptr(), packed_recv_x_scales_ptr, packed_recv_src_info.data_ptr(), + packed_recv_layout_range.data_ptr(), packed_recv_count.data_ptr(), x.data_ptr(), + topk_idx.data_ptr(), config, current_buf, next_buf, transport, workspace, + launch_stream, phase); }; - launcher(return_recv_hook ? LOW_LATENCY_SEND_PHASE : (LOW_LATENCY_SEND_PHASE | LOW_LATENCY_RECV_PHASE)); + launcher(return_recv_hook ? low_latency::SEND_ONLY : low_latency::SEND_AND_RECV); std::optional event; if (async) { @@ -1903,22 +1934,25 @@ Buffer::low_latency_dispatch(const torch::Tensor& x, const torch::Tensor& topk_i } std::optional> recv_hook = std::nullopt; - if (return_recv_hook) recv_hook = [=]() { launcher(LOW_LATENCY_RECV_PHASE); }; + if (return_recv_hook) recv_hook = [=]() { launcher(low_latency::RECV_ONLY); }; return {packed_recv_x, packed_recv_x_scales, packed_recv_count, packed_recv_src_info, packed_recv_layout_range, event, recv_hook}; } std::tuple, std::optional>> Buffer::low_latency_combine( - const torch::Tensor& x, const torch::Tensor& topk_idx, const torch::Tensor& topk_weights, - const torch::Tensor& src_info, const torch::Tensor& layout_range, int num_max_dispatch_tokens_per_rank, - int num_experts, bool zero_copy, bool async, bool return_recv_hook, const std::optional& out) { + const torch::Tensor& x, const std::optional& x_scales, const torch::Tensor& topk_idx, + const torch::Tensor& topk_weights, const torch::Tensor& src_info, const torch::Tensor& layout_range, + int num_max_dispatch_tokens_per_rank, int num_experts, bool zero_copy, bool async, bool return_recv_hook, + const std::optional& out) { EP_HOST_ASSERT(low_latency_mode); - EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous() and x.scalar_type() == torch::kBFloat16); + const auto input_dtype = x.scalar_type(); + EP_HOST_ASSERT(input_dtype == torch::kBFloat16 || input_dtype == torch::kFloat8_e4m3fn); + EP_HOST_ASSERT(x.dim() == 3 and x.is_contiguous()); EP_HOST_ASSERT(x.size(0) == num_experts / num_ranks); EP_HOST_ASSERT(x.size(1) == num_ranks * num_max_dispatch_tokens_per_rank); - EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0); + EP_HOST_ASSERT((x.size(2) * x.element_size()) % sizeof(int4) == 0 and x.size(2) % 128 == 0); EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous()); EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and topk_idx.size(1) == topk_weights.size(1)); EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64); @@ -1931,8 +1965,23 @@ std::tuple, std::optional