From 5e0c1de254eedf96ae7c5dcaa165f70e95ce7624 Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Mon, 22 Jun 2026 17:54:46 +0000 Subject: [PATCH] WIP --- python/mscclpp/ext/ep/README.md | 8 +- python/mscclpp/ext/ep/__init__.py | 13 +- .../ext/ep/{buffer.py => communicator.py} | 85 +++--- src/ext/ep/CMakeLists.txt | 39 +-- src/ext/ep/README.md | 42 +-- src/ext/ep/bindings.cpp | 264 +++++++++++++++--- src/ext/ep/buffer.cc | 24 +- src/ext/ep/buffer.hpp | 15 +- .../python/ext/ep/test_internode_multirank.py | 16 +- .../python/ext/ep/test_intranode_multirank.py | 14 +- .../ext/ep/test_low_latency_multirank.py | 51 ++-- 11 files changed, 367 insertions(+), 204 deletions(-) rename python/mscclpp/ext/ep/{buffer.py => communicator.py} (89%) diff --git a/python/mscclpp/ext/ep/README.md b/python/mscclpp/ext/ep/README.md index 7a9a71ce..243264b8 100644 --- a/python/mscclpp/ext/ep/README.md +++ b/python/mscclpp/ext/ep/README.md @@ -70,9 +70,7 @@ class MoECommunicatorConfig: input_dtype: Optional[torch.dtype] = None quant_format: Optional[str] = None - # Scratch / transport resources - num_nvl_bytes: Optional[int] = None - num_rdma_bytes: Optional[int] = None + # Transport resources num_rdma_qps_per_rank: int = 12 # RDMA QPs per peer rank; advanced tuning num_sms: int = 20 @@ -107,7 +105,7 @@ The class should cache: | `comm` | MSCCL++ communicator or `CommGroup` | | `rank`, `world_size` | global EP rank information | | `local_rank`, `device` | CUDA device binding | -| `buffer` / `runtime` | pybind/C++ EP buffer implementation | +| internal runtime | nanobind/C++ EP runtime implementation | | `comm_stream` | optional stream for async dispatch/combine | ### Expert placement fields @@ -145,7 +143,7 @@ a later version can add an explicit `expert_map` for arbitrary placement. | `output_layout` | MLP input layout returned by dispatch | | `max_tokens_per_rank` | dispatch capacity | | `max_recv_tokens_per_rank` | recv buffer capacity | -| `num_nvl_bytes`, `num_rdma_bytes` | scratch buffer sizing | +| 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 | diff --git a/python/mscclpp/ext/ep/__init__.py b/python/mscclpp/ext/ep/__init__.py index cc79462a..45316941 100644 --- a/python/mscclpp/ext/ep/__init__.py +++ b/python/mscclpp/ext/ep/__init__.py @@ -2,31 +2,30 @@ # 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 API, and ``ExpertParallelRuntime`` is +the low-level runtime wrapper for advanced HT/intranode paths. """ -from .buffer import ( # noqa: F401 - Buffer, +from .communicator import ( # noqa: F401 CommOverlapConfig, Config, DispatchHandle, DispatchOutput, EventHandle, + ExpertParallelRuntime, MoECommunicator, MoECommunicatorConfig, QuantScales, ) __all__ = [ - "Buffer", "CommOverlapConfig", "Config", "DispatchHandle", "DispatchOutput", "EventHandle", + "ExpertParallelRuntime", "MoECommunicator", "MoECommunicatorConfig", "QuantScales", diff --git a/python/mscclpp/ext/ep/buffer.py b/python/mscclpp/ext/ep/communicator.py similarity index 89% rename from python/mscclpp/ext/ep/buffer.py rename to python/mscclpp/ext/ep/communicator.py index c98d376c..41bd548a 100644 --- a/python/mscclpp/ext/ep/buffer.py +++ b/python/mscclpp/ext/ep/communicator.py @@ -5,9 +5,10 @@ # 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. +This is a thin wrapper around the nanobind extension ``mscclpp_ep_cpp``. +``MoECommunicator`` is the high-level API. ``ExpertParallelRuntime`` is a +lower-level runtime wrapper used by legacy HT/intranode tests and advanced +callers. Current status (see ``src/ext/ep/README.md``): @@ -61,8 +62,6 @@ class MoECommunicatorConfig: output_layout: Optional[str] = None input_dtype: Optional[torch.dtype] = None quant_format: Optional[str] = None - num_nvl_bytes: Optional[int] = None - num_rdma_bytes: Optional[int] = None num_rdma_qps_per_rank: int = 12 num_sms: int = 20 comm_stream: Optional[torch.cuda.Stream] = None @@ -116,8 +115,8 @@ class CommOverlapConfig: block_ready_value: Optional[int] = None -class Buffer: - """Core expert-parallel (EP) communication buffer. +class ExpertParallelRuntime: + """Low-level expert-parallel runtime wrapper. Parameters ---------- @@ -132,7 +131,7 @@ class Buffer: low_latency_mode: Enable low-latency buffer setup for :class:`MoECommunicator`. New callers should use :class:`MoECommunicator` instead of direct LL - methods on ``Buffer``. + methods on this low-level runtime. num_qps_per_rank: Ignored for intranode mode. """ @@ -156,61 +155,64 @@ class Buffer: 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) + self._cpp_buffer = _cpp.ExpertParallelRuntime( + 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() + local_device_id = self._cpp_buffer.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() + local_ipc_handle = self._cpp_buffer.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)` + # setups because the C++ runtime 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() + root_unique_id = self._cpp_buffer.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) + self._cpp_buffer.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)) + self._cpp_buffer.sync(device_ids, ipc_handles_ba, bytearray(root_unique_id)) # ------------------------------------------------------------------ # Sanity helpers # ------------------------------------------------------------------ def is_available(self) -> bool: - return self._runtime.is_available() + return self._cpp_buffer.is_available() def is_internode_available(self) -> bool: - return self._runtime.is_internode_available() + return self._cpp_buffer.is_internode_available() def get_local_device_id(self) -> int: - return self._runtime.get_local_device_id() + return self._cpp_buffer.get_local_device_id() def get_num_rdma_ranks(self) -> int: - return self._runtime.get_num_rdma_ranks() + return self._cpp_buffer.get_num_rdma_ranks() def get_rdma_rank(self) -> int: - return self._runtime.get_rdma_rank() + return self._cpp_buffer.get_rdma_rank() def get_root_rdma_rank(self, global_: bool) -> int: - return self._runtime.get_root_rdma_rank(global_) + return self._cpp_buffer.get_root_rdma_rank(global_) # ------------------------------------------------------------------ # Layout / dispatch / combine (thin pass-through wrappers). - # Signatures mirror deep_ep.Buffer so existing test harnesses can reuse. + # These are low-level runtime APIs for compatibility tests and advanced + # callers. New MoE code should prefer MoECommunicator. # ------------------------------------------------------------------ def get_dispatch_layout( @@ -221,33 +223,32 @@ class Buffer: async_finish: bool = False, allocate_on_comm_stream: bool = False, ): - return self._runtime.get_dispatch_layout( + return self._cpp_buffer.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) + return self._cpp_buffer.intranode_dispatch(*args, **kwargs) def intranode_combine(self, *args, **kwargs): - return self._runtime.intranode_combine(*args, **kwargs) + return self._cpp_buffer.intranode_combine(*args, **kwargs) def internode_dispatch(self, *args, **kwargs): - return self._runtime.internode_dispatch(*args, **kwargs) + return self._cpp_buffer.internode_dispatch(*args, **kwargs) def internode_combine(self, *args, **kwargs): - return self._runtime.internode_combine(*args, **kwargs) + return self._cpp_buffer.internode_combine(*args, **kwargs) 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) + return self._cpp_buffer.get_local_buffer_tensor(dtype, offset, use_rdma_buffer) class MoECommunicator: """High-level MoE communicator API for dispatch/combine. - The first implementation supports the low-latency backend and keeps the - existing :class:`Buffer` API available for advanced or compatibility use. + The first implementation supports the low-latency backend. """ def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> None: @@ -310,16 +311,14 @@ class MoECommunicator: if self.quant_format not in (None, "fp8_e4m3"): raise NotImplementedError(f"unsupported low-latency quant_format: {config.quant_format}") - num_nvl_bytes = config.num_nvl_bytes or 0 - num_rdma_bytes = config.num_rdma_bytes - if num_rdma_bytes is None: - num_rdma_bytes = _get_low_latency_rdma_size_hint( - self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts - ) + 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.comm_stream = config.comm_stream if self.comm_stream is not None: - raise NotImplementedError("custom comm_stream is not wired into the low-latency Buffer yet") + raise NotImplementedError("custom comm_stream is not wired into the low-latency runtime yet") self.enable_overlap = config.enable_overlap self.num_sms = config.num_sms self._dispatch_tokens: Optional[torch.Tensor] = None @@ -328,7 +327,7 @@ class MoECommunicator: self._dispatch_layout_range: Optional[torch.Tensor] = None self._dispatch_count: Optional[torch.Tensor] = None - self.buffer = Buffer( + self._runtime = ExpertParallelRuntime( group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, @@ -336,6 +335,12 @@ class MoECommunicator: 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, @@ -349,7 +354,7 @@ class MoECommunicator: output_tensors = self._get_dispatch_output_tensors(input.device) packed_tokens, packed_scales, num_tokens_per_expert, src_info, layout_range, _event, _hook = ( - self.buffer._runtime.low_latency_dispatch( + self._runtime._cpp_buffer.low_latency_dispatch( input, topk_ids, self.max_tokens_per_rank, @@ -457,7 +462,7 @@ class MoECommunicator: raise ValueError("FP8 expert_output requires scales captured in the dispatch handle") x_scales = handle.output_scales.local - combined, _event, _hook = self.buffer._runtime.low_latency_combine( + combined, _event, _hook = self._runtime._cpp_buffer.low_latency_combine( expert_output, x_scales, handle.topk_ids, diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index 199eb63a..3fa83f54 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -1,10 +1,8 @@ # 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). +# Builds `mscclpp_ep_cpp`, a nanobind + PyTorch C++ extension that exposes the +# EP (Mixture-of-Experts dispatch/combine) runtime to Python. # # Requires: PyTorch with its CMake integration available on CMAKE_PREFIX_PATH. # Easiest invocation: @@ -13,6 +11,12 @@ # -DCMAKE_PREFIX_PATH="$(python -c 'import torch;print(torch.utils.cmake_prefix_path)')" find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) +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(Torch QUIET) if(NOT Torch_FOUND) message(WARNING @@ -22,23 +26,6 @@ if(NOT Torch_FOUND) return() 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() - endif() - endforeach() - if(NOT PYBIND11_INCLUDE_DIR) - message(FATAL_ERROR - "pybind11 not found and not bundled with Torch. Install it (pip install pybind11) " - "or provide -Dpybind11_DIR=...") - endif() -endif() - file(GLOB_RECURSE EP_SOURCES CONFIGURE_DEPENDS buffer.cc bindings.cpp @@ -55,7 +42,7 @@ 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}) target_compile_definitions(mscclpp_ep_cpp PRIVATE TORCH_EXTENSION_NAME=mscclpp_ep_cpp) target_compile_definitions(mscclpp_ep_cpp PRIVATE MSCCLPP_AVOID_BFLOAT16_ALIAS) @@ -66,18 +53,12 @@ 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 +# libtorch_python contains Python 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") diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 82171709..ea6542cb 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -13,7 +13,7 @@ targeting: | Feature | Status | |------------------------------------|---------------------------------------------| -| `Buffer` construction + IPC + sync | ✅ ported (NVLink + RDMA) | +| Runtime 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) | @@ -25,7 +25,7 @@ targeting: | 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 | +| 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). @@ -76,7 +76,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,14 +92,14 @@ 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 + internode dispatch/combine (or `cached_notify`) is tracked in the test's `XXX` comment. ## Build @@ -114,12 +114,12 @@ cmake -S . -B build \ cmake --build build -j ``` -This produces `mscclpp_ep_cpp.so` — a pybind11 PyTorch extension module. +This produces `mscclpp_ep_cpp.so` — a nanobind PyTorch 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=...) +runtime = ep.ExpertParallelRuntime(group, num_nvl_bytes=..., num_rdma_bytes=...) ``` ### Build-time CMake options @@ -166,7 +166,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,17 +268,17 @@ 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.ExpertParallelRuntime` construction. ## Layout ``` src/ext/ep/ -├── CMakeLists.txt — builds mscclpp_ep_cpp (Torch + pybind11) +├── CMakeLists.txt — builds mscclpp_ep_cpp (Torch + 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 +├── bindings.cpp — nanobind module definition └── kernels/ ├── api.cuh — host-callable kernel prototypes ├── configs.cuh — compile-time constants (GPU-only) @@ -293,8 +293,8 @@ src/ext/ep/ └── 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 MoECommunicator / ExpertParallelRuntime +└── communicator.py — torch.distributed-aware frontend test/python/ext/ep/ ├── test_intranode_multirank.py — intranode HT dispatch+combine @@ -321,7 +321,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,7 +445,7 @@ 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 @@ -454,16 +454,16 @@ Env knobs: - [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`. @@ -488,7 +488,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 +514,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 f5acfcd4..997d6e6e 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -4,82 +4,258 @@ // 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 + +#include +#include +#include +#include #include "buffer.hpp" #include "config.hpp" -namespace py = pybind11; +namespace nb = nanobind; -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +namespace nanobind::detail { + +template <> +struct type_caster { + NB_TYPE_CASTER(torch::Tensor, const_name("torch.Tensor")); + + bool from_python(handle src, uint8_t, cleanup_list*) noexcept { + bool isTensor = false; + try { + isTensor = THPVariable_Check(src.ptr()); + } catch (...) { + return false; + } + if (!isTensor) return false; + value = THPVariable_Unpack(src.ptr()); + return true; + } + + static handle from_cpp(const torch::Tensor& tensor, rv_policy, cleanup_list*) noexcept { + return handle(THPVariable_Wrap(tensor)); + } +}; + +} // namespace nanobind::detail + +namespace { + +std::string bytesToString(nb::handle data) { + Py_buffer view; + if (PyObject_GetBuffer(data.ptr(), &view, PyBUF_SIMPLE) != 0) throw nb::python_error(); + std::string result(static_cast(view.buf), static_cast(view.len)); + PyBuffer_Release(&view); + return result; +} + +nb::bytes stringToBytes(const std::string& data) { return nb::bytes(data.data(), data.size()); } + +torch::ScalarType dtypeFromPython(nb::handle dtype) { + if (!THPDtype_Check(dtype.ptr())) throw nb::type_error("expected torch.dtype"); + return reinterpret_cast(dtype.ptr())->scalar_type; +} + +std::optional optionalTensor(nb::handle obj) { + if (obj.is_none()) return std::nullopt; + return nb::cast(obj); +} + +std::vector> optionalBytesList(nb::sequence values) { + std::vector> result; + result.reserve(static_cast(nb::len(values))); + for (size_t i = 0; i < static_cast(nb::len(values)); ++i) { + nb::object item = values[i]; + if (item.is_none()) { + result.push_back(std::nullopt); + } else { + result.push_back(bytesToString(item)); + } + } + return result; +} + +std::optional optionalBytes(nb::handle value) { + if (value.is_none()) return std::nullopt; + return bytesToString(value); +} + +} // namespace + +NB_MODULE(TORCH_EXTENSION_NAME, 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) + nb::class_(m, "Config") + .def(nb::init(), nb::arg("num_sms") = 20, nb::arg("num_max_nvl_chunked_send_tokens") = 6, + nb::arg("num_max_nvl_chunked_recv_tokens") = 256, nb::arg("num_max_rdma_chunked_send_tokens") = 6, + nb::arg("num_max_rdma_chunked_recv_tokens") = 256) .def("get_nvl_buffer_size_hint", &mscclpp::ep::Config::get_nvl_buffer_size_hint) .def("get_rdma_buffer_size_hint", &mscclpp::ep::Config::get_rdma_buffer_size_hint); m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::get_low_latency_rdma_size_hint); - py::class_(m, "EventHandle") - .def(py::init<>()) + nb::class_(m, "EventHandle") + .def(nb::init<>()) .def("current_stream_wait", &mscclpp::ep::EventHandle::current_stream_wait); - // 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. - - 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")) + nb::class_(m, "ExpertParallelRuntime") + .def(nb::init(), nb::arg("rank"), nb::arg("num_ranks"), + nb::arg("num_nvl_bytes"), nb::arg("num_rdma_bytes"), nb::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("get_local_ipc_handle", + [](const mscclpp::ep::Buffer& self) { return stringToBytes(self.get_local_ipc_handle()); }) + .def("get_local_nvshmem_unique_id", + [](const mscclpp::ep::Buffer& self) { return stringToBytes(self.get_local_nvshmem_unique_id()); }) + .def("get_local_buffer_tensor", + [](const mscclpp::ep::Buffer& self, nb::handle dtype, int64_t offset, bool useRdmaBuffer) { + return self.get_local_buffer_tensor(dtypeFromPython(dtype), offset, useRdmaBuffer); + }) .def("create_unique_id", [](const mscclpp::ep::Buffer& self) { auto uid = self.create_unique_id(); - return py::bytes(reinterpret_cast(uid.data()), uid.size()); + return nb::bytes(reinterpret_cast(uid.data()), uid.size()); }) .def("connect", - [](mscclpp::ep::Buffer& self, py::bytes data) { - std::string s = data; + [](mscclpp::ep::Buffer& self, nb::handle data) { + std::string s = bytesToString(data); mscclpp::UniqueId uid; if (s.size() != uid.size()) { - throw std::runtime_error("mscclpp_ep_cpp.Buffer.connect: UniqueId size mismatch"); + throw std::runtime_error("mscclpp_ep_cpp.ExpertParallelRuntime.connect: UniqueId size mismatch"); } std::memcpy(uid.data(), s.data(), s.size()); self.connect(uid); }) - .def("sync", &mscclpp::ep::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("sync", + [](mscclpp::ep::Buffer& self, const std::vector& deviceIds, nb::sequence allGatheredHandles, + nb::handle rootUniqueId) { + self.sync(deviceIds, optionalBytesList(allGatheredHandles), optionalBytes(rootUniqueId)); + }) + .def( + "get_dispatch_layout", + [](mscclpp::ep::Buffer& self, const torch::Tensor& topkIdx, int numExperts, + std::optional previousEvent, bool async, bool allocateOnCommStream) { + return self.get_dispatch_layout(topkIdx, numExperts, previousEvent, async, allocateOnCommStream); + }, + nb::arg("topk_idx"), nb::arg("num_experts"), nb::arg("previous_event").none() = nb::none(), + nb::arg("async") = false, nb::arg("allocate_on_comm_stream") = false) + .def( + "intranode_dispatch", + [](mscclpp::ep::Buffer& self, const torch::Tensor& x, nb::handle xScales, nb::handle topkIdx, + nb::handle topkWeights, nb::handle numTokensPerRank, const torch::Tensor& isTokenInRank, + nb::handle numTokensPerExpert, int cachedNumRecvTokens, nb::handle cachedRankPrefixMatrix, + nb::handle cachedChannelPrefixMatrix, int expertAlignment, const mscclpp::ep::Config& config, + std::optional previousEvent, bool async, bool allocateOnCommStream) { + return self.intranode_dispatch( + x, optionalTensor(xScales), optionalTensor(topkIdx), optionalTensor(topkWeights), + optionalTensor(numTokensPerRank), isTokenInRank, optionalTensor(numTokensPerExpert), + cachedNumRecvTokens, optionalTensor(cachedRankPrefixMatrix), optionalTensor(cachedChannelPrefixMatrix), + expertAlignment, config, previousEvent, async, allocateOnCommStream); + }, + nb::arg("x"), nb::arg("x_scales").none(), nb::arg("topk_idx").none(), nb::arg("topk_weights").none(), + nb::arg("num_tokens_per_rank").none(), nb::arg("is_token_in_rank"), nb::arg("num_tokens_per_expert").none(), + nb::arg("cached_num_recv_tokens"), nb::arg("cached_rank_prefix_matrix").none(), + nb::arg("cached_channel_prefix_matrix").none(), nb::arg("expert_alignment"), nb::arg("config"), + nb::arg("previous_event").none(), nb::arg("async"), nb::arg("allocate_on_comm_stream")) + .def( + "intranode_combine", + [](mscclpp::ep::Buffer& self, const torch::Tensor& x, nb::handle topkWeights, const torch::Tensor& srcIdx, + const torch::Tensor& rankPrefixMatrix, const torch::Tensor& channelPrefixMatrix, + const torch::Tensor& sendHead, const mscclpp::ep::Config& config, + std::optional previousEvent, bool async, bool allocateOnCommStream) { + return self.intranode_combine(x, optionalTensor(topkWeights), srcIdx, rankPrefixMatrix, channelPrefixMatrix, + sendHead, config, previousEvent, async, allocateOnCommStream); + }, + nb::arg("x"), nb::arg("topk_weights").none(), nb::arg("src_idx"), nb::arg("rank_prefix_matrix"), + nb::arg("channel_prefix_matrix"), nb::arg("send_head"), nb::arg("config"), nb::arg("previous_event").none(), + nb::arg("async"), nb::arg("allocate_on_comm_stream")) + .def( + "internode_dispatch", + [](mscclpp::ep::Buffer& self, const torch::Tensor& x, nb::handle xScales, nb::handle topkIdx, + nb::handle topkWeights, nb::handle numTokensPerRank, nb::handle numTokensPerRdmaRank, + const torch::Tensor& isTokenInRank, nb::handle numTokensPerExpert, int cachedNumRecvTokens, + int cachedNumRdmaRecvTokens, nb::handle cachedRdmaChannelPrefixMatrix, + nb::handle cachedRecvRdmaRankPrefixSum, nb::handle cachedGblChannelPrefixMatrix, + nb::handle cachedRecvGblRankPrefixSum, int expertAlignment, const mscclpp::ep::Config& config, + std::optional previousEvent, bool async, bool allocateOnCommStream) { + return self.internode_dispatch( + x, optionalTensor(xScales), optionalTensor(topkIdx), optionalTensor(topkWeights), + optionalTensor(numTokensPerRank), optionalTensor(numTokensPerRdmaRank), isTokenInRank, + optionalTensor(numTokensPerExpert), cachedNumRecvTokens, cachedNumRdmaRecvTokens, + optionalTensor(cachedRdmaChannelPrefixMatrix), optionalTensor(cachedRecvRdmaRankPrefixSum), + optionalTensor(cachedGblChannelPrefixMatrix), optionalTensor(cachedRecvGblRankPrefixSum), + expertAlignment, config, previousEvent, async, allocateOnCommStream); + }, + nb::arg("x"), nb::arg("x_scales").none(), nb::arg("topk_idx").none(), nb::arg("topk_weights").none(), + nb::arg("num_tokens_per_rank").none(), nb::arg("num_tokens_per_rdma_rank").none(), + nb::arg("is_token_in_rank"), nb::arg("num_tokens_per_expert").none(), nb::arg("cached_num_recv_tokens"), + nb::arg("cached_num_rdma_recv_tokens"), nb::arg("cached_rdma_channel_prefix_matrix").none(), + nb::arg("cached_recv_rdma_rank_prefix_sum").none(), nb::arg("cached_gbl_channel_prefix_matrix").none(), + nb::arg("cached_recv_gbl_rank_prefix_sum").none(), nb::arg("expert_alignment"), nb::arg("config"), + nb::arg("previous_event").none(), nb::arg("async"), nb::arg("allocate_on_comm_stream")) + .def( + "internode_combine", + [](mscclpp::ep::Buffer& self, const torch::Tensor& x, nb::handle topkWeights, const torch::Tensor& srcMeta, + const torch::Tensor& isCombinedTokenInRank, const torch::Tensor& rdmaChannelPrefixMatrix, + const torch::Tensor& rdmaRankPrefixSum, const torch::Tensor& gblChannelPrefixMatrix, + const torch::Tensor& combinedRdmaHead, const torch::Tensor& combinedNvlHead, + const mscclpp::ep::Config& config, std::optional previousEvent, bool async, + bool allocateOnCommStream) { + return self.internode_combine(x, optionalTensor(topkWeights), srcMeta, isCombinedTokenInRank, + rdmaChannelPrefixMatrix, rdmaRankPrefixSum, gblChannelPrefixMatrix, + combinedRdmaHead, combinedNvlHead, config, previousEvent, async, + allocateOnCommStream); + }, + nb::arg("x"), nb::arg("topk_weights").none(), nb::arg("src_meta"), nb::arg("is_combined_token_in_rank"), + nb::arg("rdma_channel_prefix_matrix"), nb::arg("rdma_rank_prefix_sum"), nb::arg("gbl_channel_prefix_matrix"), + nb::arg("combined_rdma_head"), nb::arg("combined_nvl_head"), nb::arg("config"), + nb::arg("previous_event").none(), nb::arg("async"), nb::arg("allocate_on_comm_stream")) .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, py::arg("x"), py::arg("x_scales"), - py::arg("topk_idx"), py::arg("topk_weights"), py::arg("src_info"), py::arg("layout_range"), - py::arg("num_max_dispatch_tokens_per_rank"), py::arg("num_experts"), py::arg("zero_copy"), py::arg("async"), - py::arg("return_recv_hook"), py::arg("out") = py::none()) + .def( + "low_latency_dispatch", + [](mscclpp::ep::Buffer& self, const torch::Tensor& x, const torch::Tensor& topkIdx, + int numMaxDispatchTokensPerRank, int numExperts, bool useFp8, bool async, bool returnRecvHook, + nb::handle outPackedRecvX, nb::handle outPackedRecvXScales, nb::handle outPackedRecvSrcInfo, + nb::handle outPackedRecvLayoutRange, nb::handle outPackedRecvCount) { + return self.low_latency_dispatch(x, topkIdx, numMaxDispatchTokensPerRank, numExperts, useFp8, async, + returnRecvHook, optionalTensor(outPackedRecvX), + optionalTensor(outPackedRecvXScales), optionalTensor(outPackedRecvSrcInfo), + optionalTensor(outPackedRecvLayoutRange), + optionalTensor(outPackedRecvCount)); + }, + nb::arg("x"), nb::arg("topk_idx"), nb::arg("num_max_dispatch_tokens_per_rank"), nb::arg("num_experts"), + nb::arg("use_fp8"), nb::arg("async"), nb::arg("return_recv_hook"), + nb::arg("out_packed_recv_x").none() = nb::none(), nb::arg("out_packed_recv_x_scales").none() = nb::none(), + nb::arg("out_packed_recv_src_info").none() = nb::none(), + nb::arg("out_packed_recv_layout_range").none() = nb::none(), + nb::arg("out_packed_recv_count").none() = nb::none()) + .def( + "low_latency_combine", + [](mscclpp::ep::Buffer& self, const torch::Tensor& x, nb::handle xScales, const torch::Tensor& topkIdx, + const torch::Tensor& topkWeights, const torch::Tensor& srcInfo, const torch::Tensor& layoutRange, + int numMaxDispatchTokensPerRank, int numExperts, bool zeroCopy, bool async, bool returnRecvHook, + nb::handle out) { + return self.low_latency_combine(x, optionalTensor(xScales), topkIdx, topkWeights, srcInfo, layoutRange, + numMaxDispatchTokensPerRank, numExperts, zeroCopy, async, returnRecvHook, + optionalTensor(out)); + }, + nb::arg("x"), nb::arg("x_scales").none(), nb::arg("topk_idx"), nb::arg("topk_weights"), nb::arg("src_info"), + nb::arg("layout_range"), nb::arg("num_max_dispatch_tokens_per_rank"), nb::arg("num_experts"), + nb::arg("zero_copy"), nb::arg("async"), nb::arg("return_recv_hook"), nb::arg("out").none() = nb::none()) .def("get_next_low_latency_combine_buffer", &mscclpp::ep::Buffer::get_next_low_latency_combine_buffer); } diff --git a/src/ext/ep/buffer.cc b/src/ext/ep/buffer.cc index 64090639..91a9778a 100644 --- a/src/ext/ep/buffer.cc +++ b/src/ext/ep/buffer.cc @@ -3,8 +3,6 @@ #include #include #include -#include -#include #include #include @@ -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); @@ -1323,7 +1316,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); diff --git a/src/ext/ep/buffer.hpp b/src/ext/ep/buffer.hpp index 70dccb5a..44d642ae 100644 --- a/src/ext/ep/buffer.hpp +++ b/src/ext/ep/buffer.hpp @@ -2,8 +2,6 @@ // Licensed under the MIT License. #pragma once -#include -#include #include #include @@ -11,6 +9,8 @@ #include #include #include +#include +#include #include #include @@ -207,19 +207,18 @@ struct Buffer { int get_local_device_id() const; - pybind11::bytearray get_local_ipc_handle() const; + std::string get_local_ipc_handle() const; - pybind11::bytearray get_local_nvshmem_unique_id() const; + std::string get_local_nvshmem_unique_id() const; - torch::Tensor get_local_buffer_tensor(const pybind11::object& dtype, int64_t offset, bool use_rdma_buffer) const; + torch::Tensor get_local_buffer_tensor(torch::ScalarType dtype, int64_t offset, bool use_rdma_buffer) const; mscclpp::UniqueId create_unique_id() const; void connect(mscclpp::UniqueId root_id); - void sync(const std::vector& device_ids, - const std::vector>& all_gathered_handles, - const std::optional& root_unique_id_opt); + void sync(const std::vector& device_ids, const std::vector>& all_gathered_handles, + const std::optional& root_unique_id_opt); std::tuple, torch::Tensor, torch::Tensor, std::optional> get_dispatch_layout(const torch::Tensor& topk_idx, int num_experts, std::optional& previous_event, diff --git a/test/python/ext/ep/test_internode_multirank.py b/test/python/ext/ep/test_internode_multirank.py index 29d2c49f..c38b0945 100644 --- a/test/python/ext/ep/test_internode_multirank.py +++ b/test/python/ext/ep/test_internode_multirank.py @@ -139,7 +139,7 @@ def main(): x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - # Buffer config for internode HT: needs num_rdma_bytes > 0. Size buffers + # Runtime config for internode HT: needs num_rdma_bytes > 0. Size buffers # using max(hidden, bench_hidden) so the optional bench phase fits. cfg = ep.Config( int(os.environ.get("MSCCLPP_EP_NSM", "152")), @@ -160,10 +160,12 @@ def main(): flush=True, ) - print(f"[rank {rank}] creating Buffer", flush=True) - buf = ep.Buffer(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False) + print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True) + buf = ep.ExpertParallelRuntime( + group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False + ) print( - f"[rank {rank}] Buffer created is_available={buf.is_available()} " + f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()} " f"is_internode={buf.is_internode_available()}", flush=True, ) @@ -248,7 +250,7 @@ def main(): # the various *_channel_prefix_matrix tensors can still be in flight on # the comm stream when combine launches, producing a deadlock inside the # combine forwarder (NVL check never advances). Investigate proper - # stream-dependency hand-off in Buffer::internode_dispatch. + # stream-dependency hand-off in ExpertParallelRuntime.internode_dispatch. torch.cuda.synchronize() dist.barrier(group=group) @@ -319,7 +321,7 @@ def main(): print(f"[bench] skip: topk={bench_num_topk} > experts={bench_num_experts}", flush=True) return - # Respect the Buffer's pre-sized num_nvl_bytes / num_rdma_bytes budget. + # Respect the runtime's pre-sized num_nvl_bytes / num_rdma_bytes budget. per_peer_nvl = num_nvl_bytes // max(1, num_ranks) per_peer_rdma = num_rdma_bytes // max(1, num_ranks) if bench_hidden * x.element_size() > min(per_peer_nvl, per_peer_rdma): @@ -327,7 +329,7 @@ def main(): print( f"[bench] skip: hidden={bench_hidden} bytes/row={bench_hidden * x.element_size()} " f">= min(per-peer NVL {per_peer_nvl}, RDMA {per_peer_rdma}). " - f"Rerun with a larger Buffer or smaller hidden.", + f"Rerun with a larger runtime or smaller hidden.", flush=True, ) return diff --git a/test/python/ext/ep/test_intranode_multirank.py b/test/python/ext/ep/test_intranode_multirank.py index 8ca757d3..5860aeb0 100644 --- a/test/python/ext/ep/test_intranode_multirank.py +++ b/test/python/ext/ep/test_intranode_multirank.py @@ -5,7 +5,7 @@ Launch with: torchrun --nproc_per_node= test/python/ext/ep/test_intranode_multirank.py -Tests that Buffer::sync() succeeds across N GPUs on a single node and that +Tests that ExpertParallelRuntime sync succeeds across N GPUs on a single node and that a round-trip dispatch + combine preserves data (sum of top-k weighted copies). Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass @@ -104,7 +104,7 @@ def main(): # Token payload = rank id (cast to bf16) so we can check correctness x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - # Allocate Buffer (intranode only: num_rdma_bytes=0). Size the NVL buffer + # Allocate runtime (intranode only: num_rdma_bytes=0). Size the NVL buffer # using max(hidden, bench_hidden) so the optional bench phase fits. cfg = ep.Config( int(os.environ.get("MSCCLPP_EP_NUM_SMS", "20")), @@ -121,9 +121,9 @@ def main(): flush=True, ) - print(f"[rank {rank}] creating Buffer", flush=True) - buf = ep.Buffer(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False) - print(f"[rank {rank}] Buffer created is_available={buf.is_available()}", flush=True) + print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True) + buf = ep.ExpertParallelRuntime(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False) + print(f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()}", flush=True) assert buf.is_available() # get_dispatch_layout sanity @@ -246,14 +246,14 @@ def main(): return # Rebuild inputs at bench size. Keep same layout recipe as above but at - # larger (num_tokens, hidden); Buffer is sized off the original cfg+hidden, + # larger (num_tokens, hidden); runtime is sized off the original cfg+hidden, # so bench must fit within num_nvl_bytes. If it doesn't, we skip. if bench_hidden * x.element_size() > (num_nvl_bytes // max(1, num_ranks)): if rank == 0: print( f"[bench] skip: hidden={bench_hidden} bytes/row={bench_hidden * x.element_size()} " f"> per-peer budget {num_nvl_bytes // num_ranks}. " - f"Rerun with a larger Buffer or smaller hidden.", + f"Rerun with a larger runtime or smaller hidden.", flush=True, ) return diff --git a/test/python/ext/ep/test_low_latency_multirank.py b/test/python/ext/ep/test_low_latency_multirank.py index e4b375ea..26e09038 100644 --- a/test/python/ext/ep/test_low_latency_multirank.py +++ b/test/python/ext/ep/test_low_latency_multirank.py @@ -3,7 +3,8 @@ """Multi-rank low-latency functional test for mscclpp_ep. Launch with (intra-node, 8 GPUs): - torchrun --nproc_per_node=8 test/python/ext/ep/test_low_latency_multirank.py + torchrun --nproc_per_node=8 test/python/ext/ep/test_low_latency_multirank.py \ + --num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256 Launch with (2 nodes, 1 GPU per node -- DeepEP's recommended LL topology): # node 0: @@ -33,9 +34,9 @@ we need for an LL port smoke test. BF16-only (no FP8 check). from __future__ import annotations +import argparse import os import random -import sys # Disable ProcessGroupNCCL's HeartbeatMonitor before importing torch.distributed. # It runs in a background thread polling the TCPStore; under mpirun, rank 0 @@ -47,6 +48,19 @@ import torch import torch.distributed as dist +def parse_args(): + parser = argparse.ArgumentParser(description="MSCCL++ EP low-latency multi-rank correctness/benchmark test") + parser.add_argument("--num-tokens", type=int, default=128) + parser.add_argument("--hidden", type=int, default=7168, help="LL kernels are compiled for a fixed hidden set") + parser.add_argument("--num-topk", type=int, default=8) + parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness") + parser.add_argument("--bench-warmup", type=int, default=5) + parser.add_argument("--bench-iters", type=int, default=20) + parser.add_argument("--local-rank", "--local_rank", type=int, default=None, help=argparse.SUPPRESS) + return parser.parse_args() + + def init_dist(): rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) @@ -62,6 +76,7 @@ def init_dist(): def main(): + args = parse_args() rank, num_ranks, local_rank, group = init_dist() from mscclpp.ext import ep @@ -69,12 +84,10 @@ def main(): rank_offset = 128 assert num_ranks - rank_offset < 257, "too many ranks for bf16 precision anchor" - num_tokens = int(os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "128")) - hidden = int( - os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168") - ) # LL kernels are compiled for a fixed set; see SWITCH_HIDDEN - num_topk = int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", "8")) - num_experts = int(os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", "256")) + num_tokens = args.num_tokens + hidden = args.hidden + num_topk = args.num_topk + num_experts = args.num_experts assert num_experts % num_ranks == 0 num_local_experts = num_experts // num_ranks @@ -103,20 +116,18 @@ def main(): mode="ll", num_rdma_qps_per_rank=max(1, num_experts // num_ranks), ) - buf = moe_comm.buffer if rank == 0: print( f"[cfg] num_ranks={num_ranks} num_tokens={num_tokens} hidden={hidden} " - f"num_experts={num_experts} num_topk={num_topk} " - f"num_rdma_bytes={buf.num_rdma_bytes}", + f"num_experts={num_experts} num_topk={num_topk}", flush=True, ) print( - f"[rank {rank}] Buffer created is_available={buf.is_available()} " - f"is_internode={buf.is_internode_available()}", + f"[rank {rank}] MoECommunicator created is_available={moe_comm.is_available()} " + f"is_internode={moe_comm.is_internode_available()}", flush=True, ) - assert buf.is_available() + assert moe_comm.is_available() dist.barrier(group=group) torch.cuda.synchronize() @@ -188,15 +199,15 @@ def main(): print("PASS", flush=True) # ------------------------------------------------------------------ - # Optional benchmark (enable with MSCCLPP_EP_BENCH=1). Times dispatch - # and combine separately, reporting per-iter latency (max across ranks) - # and aggregate effective bandwidth (sum across ranks). + # Optional benchmark. Times dispatch and combine separately, reporting + # per-iter latency (max across ranks) and aggregate effective bandwidth + # (sum across ranks). # ------------------------------------------------------------------ - if os.environ.get("MSCCLPP_EP_BENCH", "0") != "1": + if not args.bench: return - warmup = int(os.environ.get("MSCCLPP_EP_BENCH_WARMUP", "5")) - iters = int(os.environ.get("MSCCLPP_EP_BENCH_ITERS", "20")) + warmup = args.bench_warmup + iters = args.bench_iters def _dispatch(): return moe_comm.dispatch(x, topk_idx, topk_weights)