From 8e34326d7a9b03ca41fdfda0be2b4879754d49a3 Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Mon, 6 Jul 2026 21:14:29 -0700 Subject: [PATCH] Binyli/ep revise (#828) This pull request makes significant improvements to the MoE (Mixture of Experts) Python API and documentation, focusing on clarifying and expanding the Expert Parallel (EP) interface, especially around quantization, dispatch/combine handles, and overlap configuration. The changes introduce new data structures, update function signatures, and improve documentation to better reflect the current and planned capabilities of the system. Additionally, the base development container is updated to CUDA 13.0, and minor corrections are made to extension naming. --- .devcontainer/devcontainer.json | 2 +- docs/quickstart.md | 2 +- python/csrc/core_py.cpp | 8 +- python/mscclpp/{ext => }/ep/README.md | 242 +++-- python/mscclpp/ep/__init__.py | 54 ++ python/mscclpp/ep/_cpp.py | 23 + python/mscclpp/ep/communicator.py | 170 ++++ .../ep/buffer.py => ep/high_throughput.py} | 470 +++++++-- python/mscclpp/ep/low_latency.py | 332 +++++++ python/mscclpp/ep/types.py | 205 ++++ python/mscclpp/ep/utils.py | 148 +++ python/mscclpp/ext/__init__.py | 6 - python/mscclpp/ext/ep/__init__.py | 38 - python/mscclpp/ext/ep/communicator.py | 895 ------------------ src/ext/ep/README.md | 22 +- src/ext/ep/ht_runtime.cc | 2 +- .../{ext => }/ep/test_internode_multirank.py | 175 +--- .../{ext => }/ep/test_intranode_multirank.py | 133 +-- .../ep/test_low_latency_multirank.py | 16 +- 19 files changed, 1585 insertions(+), 1358 deletions(-) rename python/mscclpp/{ext => }/ep/README.md (76%) create mode 100644 python/mscclpp/ep/__init__.py create mode 100644 python/mscclpp/ep/_cpp.py create mode 100644 python/mscclpp/ep/communicator.py rename python/mscclpp/{ext/ep/buffer.py => ep/high_throughput.py} (58%) create mode 100644 python/mscclpp/ep/low_latency.py create mode 100644 python/mscclpp/ep/types.py create mode 100644 python/mscclpp/ep/utils.py delete mode 100644 python/mscclpp/ext/ep/__init__.py delete mode 100644 python/mscclpp/ext/ep/communicator.py rename test/python/{ext => }/ep/test_internode_multirank.py (75%) rename test/python/{ext => }/ep/test_intranode_multirank.py (79%) rename test/python/{ext => }/ep/test_low_latency_multirank.py (95%) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 6fc7dd32..a5d6bf23 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -3,7 +3,7 @@ "build": { "dockerfile": "Dockerfile", "args": { - "BASE_IMAGE": "ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9", + "BASE_IMAGE": "ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0", "USERNAME": "devuser", "SSH_PORT": "22345" } diff --git a/docs/quickstart.md b/docs/quickstart.md index 6230723f..f9b3d0a7 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -118,7 +118,7 @@ $ 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 +- **`ep`**: Build and install the Expert Parallel extension (`mscclpp.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). diff --git a/python/csrc/core_py.cpp b/python/csrc/core_py.cpp index a94f9863..b6f4c1c9 100644 --- a/python/csrc/core_py.cpp +++ b/python/csrc/core_py.cpp @@ -62,20 +62,20 @@ void register_core(nb::module_& m) { void* data = reinterpret_cast(ptr); self->send(data, size, peer, tag); }, - nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag")) + nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"), nb::call_guard()) .def( "recv", [](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) { void* data = reinterpret_cast(ptr); self->recv(data, size, peer, tag); }, - nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag")) + nb::arg("data"), nb::arg("size"), nb::arg("peer"), nb::arg("tag"), nb::call_guard()) .def("all_gather", &Bootstrap::allGather, nb::arg("allData"), nb::arg("size")) .def("barrier", &Bootstrap::barrier) .def("send", static_cast&, int, int)>(&Bootstrap::send), - nb::arg("data"), nb::arg("peer"), nb::arg("tag")) + nb::arg("data"), nb::arg("peer"), nb::arg("tag"), nb::call_guard()) .def("recv", static_cast&, int, int)>(&Bootstrap::recv), nb::arg("data"), - nb::arg("peer"), nb::arg("tag")); + nb::arg("peer"), nb::arg("tag"), nb::call_guard()); nb::class_(m, "CppUniqueId") .def(nb::init<>()) diff --git a/python/mscclpp/ext/ep/README.md b/python/mscclpp/ep/README.md similarity index 76% rename from python/mscclpp/ext/ep/README.md rename to python/mscclpp/ep/README.md index 40c066c2..255f9d80 100644 --- a/python/mscclpp/ext/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -30,7 +30,7 @@ The dispatch output should make the local MLP contract explicit: Use `MoECommunicator` as the public class name: ```python -from mscclpp.ext.ep import MoECommunicator +from mscclpp.ep import MoECommunicator moe_comm = MoECommunicator(...) ``` @@ -67,8 +67,7 @@ class MoECommunicatorConfig: output_layout: Optional[DispatchLayout] = None # default is derived from mode # Quantization defaults - input_dtype: Optional[torch.dtype] = None - quant_format: Optional[str] = None + quant: Optional[QuantConfig] = None # Transport resources num_rdma_qps_per_rank: int = 12 # RDMA QPs per peer rank; advanced tuning @@ -137,7 +136,7 @@ a later version can add an explicit `expert_map` for arbitrary placement. | Field | Purpose | |---|---| -| `mode` | Backend selection (`"ll"` active; `"ht"` archived/not compiled) | +| `mode` | Backend selection (`MoEMode.LOW_LATENCY` or `MoEMode.HIGH_THROUGHPUT`) | | `output_layout` | MLP input layout returned by dispatch | | `max_tokens_per_rank` | dispatch capacity | | `max_recv_tokens_per_rank` | recv buffer capacity | @@ -151,10 +150,10 @@ 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`. +The active implementation supports `mode=MoEMode.LOW_LATENCY` and +`mode=MoEMode.HIGH_THROUGHPUT`. `mode` must be a `MoEMode` enum value, not a +string. LL uses an expert-major output layout; HT uses a flat output layout and +selects intranode vs internode transport from the runtime size hints. ```python moe_comm = MoECommunicator(..., mode=MoEMode.LOW_LATENCY) @@ -190,7 +189,7 @@ class MoECommunicator: input: torch.Tensor, topk_ids: torch.Tensor, weights: Optional[torch.Tensor] = None, - scales: Optional[QuantScales] = None, + quant: Optional[QuantConfig] = None, *, output_buffer: torch.Tensor, stream: Optional[torch.cuda.Stream] = None, @@ -235,14 +234,14 @@ dispatch_out, handle = moe_comm.dispatch( input, topk_ids, weights=None, - scales=None, + quant=None, output_buffer=output_buffer, ) expert_output = mlp( dispatch_out.tokens, - dispatch_out.num_tokens_per_expert, - dispatch_out.scales, + dispatch_out.layout, + dispatch_out.quant, ) output = moe_comm.combine(expert_output, handle) @@ -251,14 +250,20 @@ 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. +`DispatchOutput.layout` carries both the layout kind (`FLAT` or `EXPERT_MAJOR`) +and layout-specific metadata. +Expert-grouped layouts populate +`num_tokens_per_expert`; future layouts that do not expose per-expert grouping +can leave those fields as `None`. + ## Proposed types ```python @dataclass -class QuantScales: - local: Optional[torch.Tensor] = None +class QuantConfig: + dtype: Optional[torch.dtype] = None + block_scales: Optional[torch.Tensor] = None global_scale: Optional[torch.Tensor] = None - format: Optional[str] = None block_size: Optional[int] = None @@ -267,34 +272,99 @@ class DispatchLayout(str, Enum): EXPERT_MAJOR = "expert_major" +@dataclass +class DispatchLayoutInfo: + kind: DispatchLayout + num_tokens_per_expert: Optional[torch.Tensor | list[int]] = None + offsets: Optional[torch.Tensor] = None + + +@dataclass +class DispatchOutputInfo: + layout: DispatchLayoutInfo + quant: Optional[QuantConfig] = 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 + quant: Optional[QuantConfig] + layout: DispatchLayoutInfo + + +@dataclass +class ExpertMajorCombineContext: + topk_ids: torch.Tensor + weights: torch.Tensor + num_experts: int + num_tokens: int + hidden_size: int + src_info: torch.Tensor + layout_range: torch.Tensor + num_max_dispatch_tokens_per_rank: int + + +@dataclass +class RowMajorIntranodeCombineContext: + ... + + +@dataclass +class RowMajorInternodeCombineContext: + ... + + +CombineContext = ExpertMajorCombineContext | RowMajorIntranodeCombineContext | RowMajorInternodeCombineContext class DispatchHandle: - """Opaque handle returned by dispatch and consumed by combine.""" + """Base opaque handle returned by dispatch and consumed by combine.""" + + output_info: DispatchOutputInfo + + +class ExpertMajorDispatchHandle(DispatchHandle): + combine_context: ExpertMajorCombineContext + + +class RowMajorIntranodeDispatchHandle(DispatchHandle): + combine_context: RowMajorIntranodeCombineContext + + +class RowMajorInternodeDispatchHandle(DispatchHandle): + combine_context: RowMajorInternodeCombineContext + + +@dataclass +class OperationOverlapConfig: + stream: Optional[torch.cuda.Stream] = None + wait_event: Optional[torch.cuda.Event] = None + num_comm_sms: Optional[int] = None + + +@dataclass +class BlockOverlapConfig: + block_size_m: int + ready_signal: torch.Tensor + ready_value: int = 1 + stream: Optional[torch.cuda.Stream] = None + wait_event: Optional[torch.cuda.Event] = None + num_comm_sms: Optional[int] = None @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 + operation: Optional[OperationOverlapConfig] = None + block: Optional[BlockOverlapConfig] = None + + @property + def level(self) -> str: ... ``` `create_overlap_config` creates optional overlap configuration for async -dispatch/combine calls. +dispatch/combine calls. The `op` argument is used only to validate construction; +the returned config describes how to overlap, not which operation will consume it. ```python dispatch_overlap_config = moe_comm.create_overlap_config(op="dispatch") @@ -320,27 +390,38 @@ combine_overlap_config = moe_comm.create_overlap_config( `op="dispatch", level="block"` is not part of the first version. Dispatch overlap is operation-level only. -`CommOverlapConfig` fields: +`CommOverlapConfig` contains exactly one overlap mode: + +| Field | Purpose | +|---|---| +| `operation` | Operation-level stream/event/SM config | +| `block` | Block-level ready-signal config | + +`OperationOverlapConfig` 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: +`BlockOverlapConfig` fields: -- 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. +| Field | Purpose | +|---|---| +| `block_size_m` | Rows/tokens per ready block | +| `ready_signal` | Device tensor written by MLP and waited on by combine | +| `ready_value` | Signal value that marks one block as ready for combine | +| `stream` | Optional communication stream | +| `wait_event` | Optional event the communication op waits on before starting | +| `num_comm_sms` | Optional SM budget for communication | + +Each concrete `DispatchHandle` stores a layout-specific `combine_context` used +to reverse dispatch and finish combine. `ExpertMajorDispatchHandle` uses +`ExpertMajorCombineContext` (`topk_ids`, `weights`, source info, layout ranges, +shape, and capacity). Row-major handles use intranode or internode combine contexts with +receive-side weights, source indices, prefix matrices, and send-head tensors. +The MLP should treat the handle as opaque and pass it back to `combine`. ## Dispatch inputs @@ -389,19 +470,20 @@ 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` +### `quant` -`scales` contains activation quantization metadata for `input`. It should be -`None` for BF16/FP16 input. +`quant` contains activation quantization metadata for `input`. It should be +`None` for BF16/FP16 input. The quantized tensor dtype is stored in +`quant.dtype`. 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 | +| Format | `input` | `quant.dtype` | `quant.block_scales` | `quant.global_scale` | +|---|---|---|---|---| +| BF16/FP16 | `[T, H]` | `None` | `None` | `None` | +| FP8 E4M3 | `[T, H]` FP8 | `torch.float8_e4m3fn` | `[T, H / block_size]`, often block size 128 | usually `None` | +| NVFP4 | backend-defined packed/logical `[T, H]` | backend-defined | block scale tensor | optional global scale | +| MXFP8 | backend-defined `[T, H]` | backend-defined | 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. @@ -421,8 +503,8 @@ 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. +For FP8 dispatch it is FP8 and the returned `DispatchOutput.quant` carries the +matching dtype and 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 @@ -450,17 +532,17 @@ expert2 tokens ... ``` -`dispatch_out.num_tokens_per_expert` is ordered by local expert id: +`dispatch_out.layout.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: +For flat layout, `dispatch_out.layout.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]] +offsets = cumsum([0] + num_tokens_per_expert) +tokens[offsets[i] : offsets[i + 1]] ``` This layout is efficient for Triton or grouped GEMM kernels because it avoids @@ -481,11 +563,11 @@ local-expert-major storage viewed as 2D: 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: +For expert `i`, only the first `dispatch_out.layout.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], :] +expert_major_tokens[i, : dispatch_out.layout.num_tokens_per_expert[i], :] ``` The remaining slots are padding or scratch space. The MLP output must keep the @@ -493,7 +575,7 @@ same layout and slot order. ### Scale output layout -If `dispatch_out.scales` is not `None`, its local scale tensor should follow +If `dispatch_out.quant` is not `None`, its block scale tensor should follow the same packed/expert-major layout as `dispatch_out.tokens`, with the hidden dimension replaced by the scale dimension. @@ -516,8 +598,8 @@ For flat expert-major output: ```python expert_output = triton_mlp( dispatch_out.tokens, - dispatch_out.num_tokens_per_expert, - dispatch_out.scales, + dispatch_out.layout, + dispatch_out.quant, ) ``` @@ -526,8 +608,8 @@ For padded expert-major output: ```python expert_output = expert_major_mlp( dispatch_out.tokens, - dispatch_out.num_tokens_per_expert, - dispatch_out.scales, + dispatch_out.layout, + dispatch_out.quant, ) ``` @@ -577,10 +659,10 @@ dispatch_out, handle = moe_comm.dispatch( input, topk_ids, weights, - scales, + quant, output_buffer=output_buffer, ) -expert_output = mlp(dispatch_out.tokens, dispatch_out.num_tokens_per_expert) +expert_output = mlp(dispatch_out.tokens, dispatch_out.layout) output = moe_comm.combine(expert_output, handle) ``` @@ -603,7 +685,7 @@ dispatch_req = moe_comm.dispatch_async( input, topk_ids, weights, - scales, + quant, output_buffer=output_buffer, overlap_config=dispatch_overlap_config, ) @@ -611,7 +693,7 @@ dispatch_req = moe_comm.dispatch_async( # 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) +expert_output = mlp(dispatch_out.tokens, dispatch_out.layout) combine_overlap_config = moe_comm.create_overlap_config(op="combine", handle=handle) combine_req = moe_comm.combine_async( @@ -646,7 +728,7 @@ combine_overlap_config = moe_comm.create_overlap_config( config = combine_overlap_config expert_output = mlp( dispatch_out.tokens, - dispatch_out.num_tokens_per_expert, + dispatch_out.layout, config=config, ) @@ -667,8 +749,8 @@ 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`. +- signal at the block granularity defined by `overlap_config.block.block_size_m`, +- use the ready value/protocol provided by `overlap_config.block`. 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. @@ -682,8 +764,8 @@ 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. +- DeepGEMM FP8 masked down-GEMM can return block metadata such as `block_size_m` + and `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. @@ -719,13 +801,13 @@ 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 + quant=None, # BF16 path output_buffer=recv_buffer, ) expert_output = triton_grouped_mlp( recv.tokens, - recv.num_tokens_per_expert, + recv.layout, ) output = moe_comm.combine(expert_output, handle) @@ -738,9 +820,9 @@ recv, handle = moe_comm.dispatch( input=x_fp8, topk_ids=topk_ids, weights=topk_weights, - scales=QuantScales( - local=x_scales, - format="fp8_e4m3", + quant=QuantConfig( + dtype=torch.float8_e4m3fn, + block_scales=x_scales, block_size=128, ), output_buffer=recv_buffer, @@ -748,8 +830,8 @@ recv, handle = moe_comm.dispatch( expert_output = fp8_grouped_mlp( recv.tokens, - recv.scales, - recv.num_tokens_per_expert, + recv.quant, + recv.layout, ) output = moe_comm.combine(expert_output, handle) diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py new file mode 100644 index 00000000..46b93b7b --- /dev/null +++ b/python/mscclpp/ep/__init__.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""MSCCL++ Expert-Parallel + + +``MoECommunicator`` is the public API. ``mode=MoEMode.LOW_LATENCY`` runs on the +LL backend; ``mode=MoEMode.HIGH_THROUGHPUT`` runs on the HT backend (GB200 TMA +direct-gather combine + all-sender dispatch). +""" + +from .communicator import ( # noqa: F401 + BlockOverlapConfig, + CommOverlapConfig, + CombineContext, + DispatchHandle, + DispatchLayout, + DispatchLayoutInfo, + DispatchOutput, + DispatchOutputInfo, + ExpertMajorDispatchHandle, + ExpertMajorCombineContext, + MoECommunicator, + MoECommunicatorConfig, + MoEMode, + OperationOverlapConfig, + QuantConfig, + RowMajorInternodeDispatchHandle, + RowMajorInternodeCombineContext, + RowMajorIntranodeDispatchHandle, + RowMajorIntranodeCombineContext, +) + +__all__ = [ + "BlockOverlapConfig", + "CommOverlapConfig", + "CombineContext", + "DispatchHandle", + "DispatchLayout", + "DispatchLayoutInfo", + "DispatchOutput", + "DispatchOutputInfo", + "ExpertMajorDispatchHandle", + "ExpertMajorCombineContext", + "MoECommunicator", + "MoECommunicatorConfig", + "MoEMode", + "OperationOverlapConfig", + "QuantConfig", + "RowMajorInternodeDispatchHandle", + "RowMajorInternodeCombineContext", + "RowMajorIntranodeDispatchHandle", + "RowMajorIntranodeCombineContext", +] diff --git a/python/mscclpp/ep/_cpp.py b/python/mscclpp/ep/_cpp.py new file mode 100644 index 00000000..dddf0038 --- /dev/null +++ b/python/mscclpp/ep/_cpp.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Shared loader for the MSCCL++ expert-parallel Python extension.""" + +from __future__ import annotations + +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 +Config = getattr(_cpp, "Config", None) + + +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/ep/communicator.py b/python/mscclpp/ep/communicator.py new file mode 100644 index 00000000..fdfa91da --- /dev/null +++ b/python/mscclpp/ep/communicator.py @@ -0,0 +1,170 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""High-level MoE dispatch/combine communicator.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from ._cpp import DispatchLayout, MoEMode +from .high_throughput import HighThroughputBackend +from .low_latency import LowLatencyBackend +from .types import ( + BlockOverlapConfig, + CommOverlapConfig, + CombineContext, + DispatchHandle, + DispatchLayoutInfo, + DispatchOutput, + DispatchOutputInfo, + ExpertMajorDispatchHandle, + ExpertMajorCombineContext, + MoECommunicatorConfig, + OperationOverlapConfig, + QuantConfig, + RowMajorInternodeDispatchHandle, + RowMajorInternodeCombineContext, + RowMajorIntranodeDispatchHandle, + RowMajorIntranodeCombineContext, +) + +__all__ = [ + "CommOverlapConfig", + "BlockOverlapConfig", + "CombineContext", + "DispatchHandle", + "DispatchLayout", + "DispatchLayoutInfo", + "DispatchOutput", + "DispatchOutputInfo", + "ExpertMajorDispatchHandle", + "ExpertMajorCombineContext", + "MoECommunicator", + "MoECommunicatorConfig", + "MoEMode", + "OperationOverlapConfig", + "QuantConfig", + "RowMajorInternodeDispatchHandle", + "RowMajorInternodeCombineContext", + "RowMajorIntranodeDispatchHandle", + "RowMajorIntranodeCombineContext", +] + + +class MoECommunicator: + """High-level MoE communicator for dispatch/combine. + + ``mode=MoEMode.LOW_LATENCY`` selects the LL backend (EXPERT_MAJOR); + ``mode=MoEMode.HIGH_THROUGHPUT`` selects the HT backend (FLAT). + """ + + 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: + config = MoECommunicatorConfig(**kwargs) + + if config.device is not None: + torch.cuda.set_device(config.device) + + if not isinstance(config.mode, MoEMode): + raise TypeError("MoECommunicatorConfig.mode must be a MoEMode") + + _validate_common_config(config) + self.mode = config.mode + self.output_layout = _resolve_output_layout(config.output_layout, self.mode) + if self.mode == MoEMode.LOW_LATENCY: + self._backend = LowLatencyBackend(config, self.output_layout) + else: + self._backend = HighThroughputBackend(config, self.output_layout) + self._publish_backend_state() + + def _publish_backend_state(self) -> None: + for name in ( + "comm", + "rank", + "world_size", + "local_rank", + "device", + "num_experts", + "hidden_size", + "topk", + "max_tokens_per_rank", + "num_sms", + "enable_overlap", + "num_local_experts", + "local_expert_start", + ): + setattr(self, name, getattr(self._backend, name)) + + def is_available(self) -> bool: + return self._backend.is_available() + + def is_internode_available(self) -> bool: + return self._backend.is_internode_available() + + def is_internode(self) -> bool: + return self._backend.is_internode() + + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor] = None, + quant: Optional[QuantConfig] = None, + *, + output_buffer: Optional[torch.Tensor] = None, + stream: Optional[torch.cuda.Stream] = None, + previous_handle: Optional[DispatchHandle] = None, + ) -> Tuple[DispatchOutput, DispatchHandle]: + return self._backend.dispatch( + input, + topk_ids, + weights, + quant, + output_buffer=output_buffer, + stream=stream, + previous_handle=previous_handle, + ) + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor] = None, + stream: Optional[torch.cuda.Stream] = None, + ) -> torch.Tensor: + return self._backend.combine(expert_output, handle, out=out, stream=stream) + + 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(operation=OperationOverlapConfig()) + + +def _validate_common_config(config: MoECommunicatorConfig) -> None: + if config.num_experts <= 0 or config.hidden_size <= 0 or config.topk <= 0 or config.max_tokens_per_rank <= 0: + raise ValueError("num_experts, hidden_size, topk, and max_tokens_per_rank must be positive") + + +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 diff --git a/python/mscclpp/ext/ep/buffer.py b/python/mscclpp/ep/high_throughput.py similarity index 58% rename from python/mscclpp/ext/ep/buffer.py rename to python/mscclpp/ep/high_throughput.py index 41469f0c..081d61e5 100644 --- a/python/mscclpp/ext/ep/buffer.py +++ b/python/mscclpp/ep/high_throughput.py @@ -3,9 +3,10 @@ # # Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), # branch ``chhwang/dev-atomic-add-cleanup``. Licensed under the MIT License. -"""Low-level HT (high-throughput) runtime wrapper for the MSCCL++ EP extension. +"""High-throughput backend for the high-level MoE communicator. -This is a thin wrapper around the nanobind extension +This module contains both the high-level HT backend used by +``MoECommunicator`` and the raw-pointer wrapper around the nanobind extension ``mscclpp_ep_cpp.ExpertParallelRuntime`` (the DeepEP-style high-throughput runtime). The extension exposes a **torch-free, raw-pointer** boundary identical in spirit to the low-latency ``MoERuntime``: every device tensor crosses the @@ -27,75 +28,43 @@ internode dispatch path: The cached fast path skips the notify phase (``cached_mode=True``) by reusing a previous dispatch's prefix matrices and recv count. -The low-latency path is served by :class:`mscclpp.ext.ep.MoERuntime`; this -runtime exposes only the HT dispatch/combine methods. +The low-latency path is served by ``low_latency.py``. """ from __future__ import annotations -from typing import List, Optional, Tuple +from typing import Any, List, Optional 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 or install via `pip install` after the build." - ) from exc - -Config = _cpp.Config +from ._cpp import Config, DispatchLayout, MoEMode, _cpp +from .types import ( + DispatchHandle, + DispatchLayoutInfo, + DispatchOutput, + DispatchOutputInfo, + RowMajorInternodeDispatchHandle, + RowMajorInternodeCombineContext, + RowMajorIntranodeDispatchHandle, + RowMajorIntranodeCombineContext, + MoECommunicatorConfig, + QuantConfig, +) +from .utils import ( + all_gather_object as _all_gather_object, + bf16_view as _bf16_view, + broadcast_object as _broadcast_object, + current_stream_ptr as _stream_ptr, + exclusive_cumsum, + ptr as _ptr, + resolve_expert_placement, +) -# ---------------------------------------------------------------------------- -# Raw-pointer helpers (the boundary is now data_ptr()-based, like MoERuntime). -# ---------------------------------------------------------------------------- - - -def _ptr(t: Optional[torch.Tensor]) -> int: - """``tensor.data_ptr()`` for a tensor, or 0 (== nullptr) for ``None``.""" - return 0 if t is None else t.data_ptr() - - -def _stream_ptr() -> int: - """Raw pointer of the current CUDA stream (matches the C++ ``cudaStream_t``).""" - return torch.cuda.current_stream().cuda_stream - - -class _DevicePointerArray: - """Minimal ``__cuda_array_interface__`` holder wrapping an existing device - pointer (no allocation, no ownership). Used to view this rank's recv pool as - a tensor for the zero-copy direct dispatch path, mirroring the old - ``torch::from_blob`` on ``recv_pool_local_ptr_``.""" - - def __init__(self, ptr: int, shape: Tuple[int, ...], typestr: str, owner) -> None: - # ``owner`` keeps the runtime (and therefore the pool allocation) alive - # for as long as the resulting tensor is referenced. - self._owner = owner - self.__cuda_array_interface__ = { - "data": (ptr, False), - "shape": shape, - "typestr": typestr, - "version": 3, - "strides": None, - } - - -def _bf16_view(ptr: int, num_tokens: int, hidden: int, owner) -> torch.Tensor: - """View a raw device pointer as a ``[num_tokens, hidden]`` bfloat16 tensor. - - bfloat16 has no ``__cuda_array_interface__`` typestr, so the memory is - imported as uint16 and reinterpreted with ``.view(torch.bfloat16)``.""" - u16 = torch.as_tensor(_DevicePointerArray(ptr, (num_tokens, hidden), " None: if low_latency_mode: raise NotImplementedError( - "ExpertParallelRuntime serves the high-throughput path only; use MoERuntime for low latency." + "HighThroughputRuntime serves the high-throughput path only; use MoERuntime for low latency." ) if num_qps_per_rank <= 0: raise ValueError("num_qps_per_rank must be > 0") - self.rank: int = group.rank() - self.group_size: int = group.size() - self.group = group + 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 @@ -128,20 +97,16 @@ class ExpertParallelRuntime: self.runtime = _cpp.ExpertParallelRuntime(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes) # Exchange device ids + CUDA-IPC handles + (for RDMA) the MSCCL++ unique id. - 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) + device_ids = _all_gather_object(comm, local_device_id, 0xE000) - 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) + ipc_handles = _all_gather_object(comm, local_ipc_handle, 0xE100) root_unique_id: Optional[bytes] = None if self.rank == 0: root_unique_id = self.runtime.create_unique_id() - broadcast_list = [root_unique_id] - dist.broadcast_object_list(broadcast_list, src=0, group=group) - root_unique_id = broadcast_list[0] + root_unique_id = _broadcast_object(comm, root_unique_id, 0, 0xE200) assert root_unique_id is not None self.runtime.connect(root_unique_id) @@ -614,16 +579,357 @@ class ExpertParallelRuntime: ) return combined_x, combined_topk_weights - # ------------------------------------------------------------------ - # 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) +class HighThroughputBackend: + """Backend implementation for ``MoEMode.HIGH_THROUGHPUT``.""" + def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) -> None: + comm = config.comm + if comm is None: + raise ValueError("mode=HIGH_THROUGHPUT requires an mscclpp.CommGroup via comm=") + if Config is None or not hasattr(_cpp, "ExpertParallelRuntime"): + raise ImportError( + "mscclpp_ep_cpp was built without the high-throughput EP backend. " + "Rebuild with -DMSCCLPP_BUILD_EXT_EP=ON and ensure Config/ExpertParallelRuntime are exported." + ) -# Backward-compatible alias for the former DeepEP-style name. -Buffer = ExpertParallelRuntime + self.comm = comm + self.rank = comm.my_rank + self.world_size = comm.nranks + self.local_rank = torch.cuda.current_device() + self.device = torch.device("cuda", self.local_rank) + self.mode = MoEMode.HIGH_THROUGHPUT + self.output_layout = output_layout + + 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 + self.num_sms = config.num_sms + self.enable_overlap = config.enable_overlap + + if self.output_layout != DispatchLayout.FLAT: + raise NotImplementedError("HT mode currently supports only DispatchLayout.FLAT") + + self.num_local_experts, self.local_expert_start = resolve_expert_placement( + num_experts=self.num_experts, + world_size=self.world_size, + rank=self.rank, + num_local_experts=config.num_local_experts, + local_expert_start=config.local_expert_start, + ) + + if config.quant is not None: + raise NotImplementedError("HT quantized dispatch (scales) is not implemented yet") + + self.expert_alignment = config.expert_alignment + self._cfg = Config( + self.num_sms, + config.nvl_chunked_send, + config.nvl_chunked_recv, + config.rdma_chunked_send, + config.rdma_chunked_recv, + ) + hidden_bytes = self.hidden_size * torch.tensor([], dtype=torch.bfloat16).element_size() + num_nvl_bytes = self._cfg.get_nvl_buffer_size_hint(hidden_bytes, self.world_size) + num_rdma_bytes = self._cfg.get_rdma_buffer_size_hint(hidden_bytes, self.world_size) + self._is_internode = num_rdma_bytes > 0 + self._runtime = HighThroughputRuntime( + comm, + num_nvl_bytes=num_nvl_bytes, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=False, + num_qps_per_rank=config.num_rdma_qps_per_rank, + ) + + def is_available(self) -> bool: + return self._runtime.is_available() + + def is_internode_available(self) -> bool: + return self._runtime.is_internode_available() + + def is_internode(self) -> bool: + return self._is_internode + + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + quant: Optional[QuantConfig], + *, + output_buffer: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + previous_handle: Optional[DispatchHandle], + ) -> tuple[DispatchOutput, DispatchHandle]: + del output_buffer + if stream is not None: + with torch.cuda.stream(stream): + return self._dispatch(input, topk_ids, weights, quant, previous_handle) + return self._dispatch(input, topk_ids, weights, quant, previous_handle) + + def _dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + quant: Optional[QuantConfig], + previous_handle: Optional[DispatchHandle], + ) -> tuple[DispatchOutput, DispatchHandle]: + self._validate_dispatch_inputs(input, topk_ids, weights, quant) + if weights is None: + weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) + + cache = getattr(previous_handle, "_dispatch_cache", None) if previous_handle is not None else None + if cache is not None: + num_tokens_per_rank = cache["num_tokens_per_rank"] + num_tokens_per_rdma_rank = cache["num_tokens_per_rdma_rank"] + num_tokens_per_expert = cache["num_tokens_per_expert"] + is_token_in_rank = cache["is_token_in_rank"] + else: + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + ) = self._runtime.get_dispatch_layout(topk_ids, self.num_experts) + + if self._is_internode: + ( + recv_x, + _recv_x_scales, + _recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + _rdma_channel_prefix_matrix, + _gbl_channel_prefix_matrix, + recv_rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix, + _recv_gbl_rank_prefix_sum, + recv_src_meta, + send_rdma_head, + send_nvl_head, + ) = self._runtime.internode_dispatch( + input, + None, + topk_ids, + weights, + num_tokens_per_rank, + num_tokens_per_rdma_rank, + is_token_in_rank, + num_tokens_per_expert, + 0, + 0, + None, + None, + None, + None, + self.expert_alignment, + self._cfg, + ) + combine_context = RowMajorInternodeCombineContext( + recv_topk_weights=recv_topk_weights, + src_meta=recv_src_meta, + is_token_in_rank=is_token_in_rank, + recv_rdma_channel_prefix_matrix=recv_rdma_channel_prefix_matrix, + recv_rdma_rank_prefix_sum=recv_rdma_rank_prefix_sum, + recv_gbl_channel_prefix_matrix=recv_gbl_channel_prefix_matrix, + send_rdma_head=send_rdma_head, + send_nvl_head=send_nvl_head, + ) + dispatch_cache = ( + cache + if cache is not None + else { + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "is_token_in_rank": is_token_in_rank, + } + ) + elif cache is not None: + ( + recv_x, + _recv_x_scales, + _recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rank_prefix_matrix, + _channel_prefix_matrix, + recv_channel_prefix_matrix, + recv_src_idx, + send_head, + ) = self._runtime.intranode_dispatch( + input, + None, + None, + None, + None, + is_token_in_rank, + None, + cache["num_recv_tokens"], + cache["rank_prefix_matrix"], + cache["channel_prefix_matrix"], + self.expert_alignment, + self._cfg, + ) + combine_context = RowMajorIntranodeCombineContext( + recv_topk_weights=recv_topk_weights, + src_idx=recv_src_idx, + rank_prefix_matrix=rank_prefix_matrix, + recv_channel_prefix_matrix=recv_channel_prefix_matrix, + send_head=send_head, + ) + dispatch_cache = cache + else: + ( + recv_x, + _recv_x_scales, + _recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rank_prefix_matrix, + channel_prefix_matrix, + recv_channel_prefix_matrix, + recv_src_idx, + send_head, + ) = self._runtime.intranode_dispatch( + input, + None, + topk_ids, + weights, + num_tokens_per_rank, + is_token_in_rank, + num_tokens_per_expert, + 0, + None, + None, + self.expert_alignment, + self._cfg, + ) + combine_context = RowMajorIntranodeCombineContext( + recv_topk_weights=recv_topk_weights, + src_idx=recv_src_idx, + rank_prefix_matrix=rank_prefix_matrix, + recv_channel_prefix_matrix=recv_channel_prefix_matrix, + send_head=send_head, + ) + dispatch_cache = { + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "is_token_in_rank": is_token_in_rank, + "rank_prefix_matrix": rank_prefix_matrix, + "channel_prefix_matrix": channel_prefix_matrix, + "num_recv_tokens": int(recv_x.size(0)), + } + + output_info = DispatchOutputInfo( + layout=DispatchLayoutInfo( + kind=DispatchLayout.FLAT, + num_tokens_per_expert=num_recv_tokens_per_expert_list, + offsets=exclusive_cumsum(num_recv_tokens_per_expert_list), + ), + quant=None, + ) + dispatch_out = DispatchOutput( + tokens=recv_x, + quant=output_info.quant, + layout=output_info.layout, + ) + handle_cls = ( + RowMajorInternodeDispatchHandle + if isinstance(combine_context, RowMajorInternodeCombineContext) + else RowMajorIntranodeDispatchHandle + ) + handle = handle_cls(output_info=output_info, combine_context=combine_context) + # The torch-free HT runtime orders its work on the caller's CUDA stream + # (no separate event handle), so there is nothing to attach here. + handle._event = None # type: ignore[attr-defined] + handle._dispatch_cache = dispatch_cache # type: ignore[attr-defined] + return dispatch_out, handle + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + ) -> torch.Tensor: + if stream is not None: + with torch.cuda.stream(stream): + return self._combine(expert_output, handle, out) + return self._combine(expert_output, handle, out) + + def _combine( + self, expert_output: torch.Tensor, handle: DispatchHandle, out: Optional[torch.Tensor] + ) -> torch.Tensor: + self._validate_combine_inputs(expert_output, handle) + if isinstance(handle, RowMajorInternodeDispatchHandle): + context = handle.combine_context + combined_x, _combined_w = self._runtime.internode_combine( + expert_output, + context.recv_topk_weights, + context.src_meta, + context.is_token_in_rank, + context.recv_rdma_channel_prefix_matrix, + context.recv_rdma_rank_prefix_sum, + context.recv_gbl_channel_prefix_matrix, + context.send_rdma_head, + context.send_nvl_head, + self._cfg, + ) + elif isinstance(handle, RowMajorIntranodeDispatchHandle): + context = handle.combine_context + combined_x, _combined_w = self._runtime.intranode_combine( + expert_output, + context.recv_topk_weights, + context.src_idx, + context.rank_prefix_matrix, + context.recv_channel_prefix_matrix, + context.send_head, + self._cfg, + ) + else: + raise ValueError("DispatchHandle does not contain row-major combine context") + if out is not None: + out.copy_(combined_x) + return out + return combined_x + + def _validate_dispatch_inputs(self, input, topk_ids, weights, quant) -> None: + if quant is not None: + raise NotImplementedError("HT 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] tensor") + if input.device.type != "cuda" or input.dtype != torch.bfloat16: + raise ValueError("HT dispatch input must be a CUDA BF16 tensor") + if input.size(1) != self.hidden_size: + raise ValueError(f"input hidden size {input.size(1)} != 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 be [input.size(0), 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") + + def _validate_combine_inputs(self, expert_output, handle) -> None: + if not isinstance(handle, (RowMajorIntranodeDispatchHandle, RowMajorInternodeDispatchHandle)): + raise TypeError("handle must be a DispatchHandle returned by dispatch") + if expert_output.dim() != 2 or not expert_output.is_contiguous(): + raise ValueError("expert_output must be a contiguous [total_recv_tokens, hidden] tensor") + if expert_output.size(1) != self.hidden_size: + raise ValueError(f"expert_output hidden size {expert_output.size(1)} != configured {self.hidden_size}") + if self._is_internode != isinstance(handle, RowMajorInternodeDispatchHandle): + raise ValueError("handle transport does not match this communicator") diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/low_latency.py new file mode 100644 index 00000000..737b3acd --- /dev/null +++ b/python/mscclpp/ep/low_latency.py @@ -0,0 +1,332 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Low-latency backend for the high-level MoE communicator.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ._cpp import DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint +from .types import ( + DispatchHandle, + DispatchLayoutInfo, + DispatchOutput, + DispatchOutputInfo, + ExpertMajorDispatchHandle, + ExpertMajorCombineContext, + MoECommunicatorConfig, + QuantConfig, +) +from .utils import cuda_stream_ptr, requires_dequantization, resolve_expert_placement + + +class LowLatencyRuntime: + """Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" + + num_sms: int = 20 + + def __init__( + self, + comm: Any, + 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") + if mode != MoEMode.LOW_LATENCY: + raise NotImplementedError("LowLatencyRuntime supports only MoEMode.LOW_LATENCY") + if num_qps_per_rank <= 0: + raise ValueError("num_qps_per_rank must be > 0") + + self.mode = mode + 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) + + 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 LowLatencyBackend: + """Backend implementation for ``MoEMode.LOW_LATENCY``.""" + + def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) -> None: + comm = config.comm + if comm is None: + raise ValueError("mode=LOW_LATENCY requires an mscclpp.CommGroup via comm=") + + self.comm = comm + self.rank = comm.my_rank + self.world_size = comm.nranks + self.local_rank = torch.cuda.current_device() + self.device = torch.device("cuda", self.local_rank) + self.mode = MoEMode.LOW_LATENCY + self.output_layout = output_layout + + 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 + self.num_sms = config.num_sms + self.enable_overlap = config.enable_overlap + + if self.output_layout != DispatchLayout.EXPERT_MAJOR: + raise NotImplementedError("low-latency mode currently supports only DispatchLayout.EXPERT_MAJOR") + if self.num_experts % self.world_size != 0: + raise ValueError("low-latency mode requires num_experts divisible by world_size") + + self.num_local_experts, self.local_expert_start = resolve_expert_placement( + num_experts=self.num_experts, + world_size=self.world_size, + rank=self.rank, + num_local_experts=config.num_local_experts, + local_expert_start=config.local_expert_start, + ) + + 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") + self.quant = config.quant + self.quant_dtype = None if self.quant is None else self.quant.dtype + if self.quant is not None and self.quant_dtype is None: + raise ValueError("quant.dtype is required when quant is provided") + if self.quant_dtype not in (None, torch.float8_e4m3fn): + raise NotImplementedError(f"unsupported low-latency quant dtype: {self.quant_dtype}") + self.dispatch_requires_quantization = self.quant_dtype is not None + + num_rdma_bytes = get_low_latency_rdma_size_hint( + self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts + ) + 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 = LowLatencyRuntime( + comm, + num_nvl_bytes=0, + num_rdma_bytes=num_rdma_bytes, + mode=self.mode, + num_qps_per_rank=config.num_rdma_qps_per_rank, + ) + # LL always uses the RDMA transport, but a single-node LL job is not + # internode topology-wise. num_rdma_ranks > 1 iff world_size spans more + # than one local NVLink domain. + self._is_internode = self._runtime.get_num_rdma_ranks() > 1 + + def is_available(self) -> bool: + return self._runtime.is_available() + + def is_internode_available(self) -> bool: + return self._runtime.is_internode_available() + + def is_internode(self) -> bool: + return self._is_internode + + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + quant: Optional[QuantConfig], + *, + output_buffer: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + previous_handle: Optional[DispatchHandle], + ) -> tuple[DispatchOutput, DispatchHandle]: + del previous_handle + self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer) + if weights is None: + weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) + + out_buf, packed_scales, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) + self._runtime.cpp_runtime.dispatch( + input.data_ptr(), + topk_ids.data_ptr(), + out_buf.data_ptr(), + 0 if packed_scales is None else packed_scales.data_ptr(), + src_info.data_ptr(), + layout_range.data_ptr(), + count.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), + ) + dispatched_quant = None + if packed_scales is not None: + dispatched_quant = QuantConfig(dtype=self.quant_dtype, block_scales=packed_scales, block_size=128) + output_info = DispatchOutputInfo( + layout=DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count), + quant=dispatched_quant, + ) + dispatch_out = DispatchOutput( + tokens=out_buf, + quant=output_info.quant, + layout=output_info.layout, + ) + handle = ExpertMajorDispatchHandle( + output_info=output_info, + combine_context=ExpertMajorCombineContext( + topk_ids=topk_ids, + weights=weights, + num_experts=self.num_experts, + num_tokens=input.size(0), + hidden_size=self.hidden_size, + src_info=src_info, + layout_range=layout_range, + num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank, + ), + ) + return dispatch_out, handle + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + ) -> torch.Tensor: + self._validate_combine_inputs(expert_output, handle, out) + if not isinstance(handle, ExpertMajorDispatchHandle): + raise ValueError("DispatchHandle does not contain expert-major combine context") + context = handle.combine_context + combine_requires_dequantization = requires_dequantization(expert_output) + x_scales = None + if combine_requires_dequantization: + if handle.output_info.quant is None or handle.output_info.quant.block_scales is None: + raise ValueError("FP8 expert_output requires scales captured in the dispatch handle") + x_scales = handle.output_info.quant.block_scales + if out is None: + out = torch.empty((context.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(), + context.topk_ids.data_ptr(), + context.weights.data_ptr(), + context.src_info.data_ptr(), + context.layout_range.data_ptr(), + out.data_ptr(), + context.num_tokens, + self.hidden_size, + context.weights.size(1), + context.num_max_dispatch_tokens_per_rank, + context.num_experts, + combine_requires_dequantization, + cuda_stream_ptr(stream), + ) + return out + + def _get_dispatch_output_tensors(self, output_buffer: 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) + return ( + output_buffer, + self._dispatch_scales, + self._dispatch_src_info, + self._dispatch_layout_range, + self._dispatch_count, + ) + + def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buffer) -> None: + if output_buffer is None: + raise ValueError("output_buffer is required for low-latency dispatch") + if quant is not None and (quant.block_scales is not None or quant.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, handle, out) -> None: + if not isinstance(handle, ExpertMajorDispatchHandle): + raise ValueError("DispatchHandle does not contain expert-major combine context") + context = handle.combine_context + if context.num_experts != self.num_experts or context.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.output_info.layout.kind == 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 = (context.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}") diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py new file mode 100644 index 00000000..2358c52d --- /dev/null +++ b/python/mscclpp/ep/types.py @@ -0,0 +1,205 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Public data types for the expert-parallel Python API.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, List, Optional, Union + +import torch +import mscclpp +from ._cpp import DispatchLayout, MoEMode + +# Quantization metadata. + + +@dataclass +class QuantConfig: + """Quantization metadata associated with an activation tensor.""" + + dtype: Optional[torch.dtype] = None + block_scales: Optional[torch.Tensor] = None + global_scale: Optional[torch.Tensor] = None + block_size: Optional[int] = None + + +# Communicator construction. + + +@dataclass +class MoECommunicatorConfig: + """Configuration for the high-level MoE dispatch/combine API.""" + + comm: Optional[mscclpp.CommGroup] = None + device: Optional[Union[torch.device, int]] = None + + # Expert topology + num_experts: int = 0 + num_local_experts: Optional[int] = None + local_expert_start: Optional[int] = None + + # 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 + + # Quantization defaults + quant: Optional[QuantConfig] = None + + # Transport / launch tuning + num_rdma_qps_per_rank: int = 12 + num_sms: int = 20 + enable_overlap: bool = False + + # HT-only buffer/launch tuning (advanced) + expert_alignment: int = 1 + nvl_chunked_send: int = 8 + nvl_chunked_recv: int = 256 + rdma_chunked_send: int = 16 + rdma_chunked_recv: int = 128 + + +# MLP-facing dispatch output. + + +@dataclass +class DispatchLayoutInfo: + """Physical layout of dispatched tokens and optional expert-group metadata.""" + + kind: DispatchLayout + num_tokens_per_expert: Optional[Union[torch.Tensor, List[int]]] = None + offsets: Optional[torch.Tensor] = None + + +@dataclass +class DispatchOutputInfo: + """Lightweight output metadata copied into both dispatch output and handle.""" + + layout: DispatchLayoutInfo + quant: Optional[QuantConfig] = None + + +@dataclass +class DispatchOutput: + """Dispatch result consumed by the local MLP.""" + + tokens: torch.Tensor + quant: Optional[QuantConfig] + layout: DispatchLayoutInfo + + +# Combine-side context. These objects are layout-specific and opaque to the MLP. + + +@dataclass +class ExpertMajorCombineContext: + """Combine context for expert-major dispatch output.""" + + topk_ids: torch.Tensor + weights: torch.Tensor + num_experts: int + num_tokens: int + hidden_size: int + src_info: torch.Tensor + layout_range: torch.Tensor + num_max_dispatch_tokens_per_rank: int + + +@dataclass +class RowMajorIntranodeCombineContext: + """Combine context for row-major intranode dispatch output.""" + + recv_topk_weights: Optional[torch.Tensor] + src_idx: torch.Tensor + rank_prefix_matrix: torch.Tensor + recv_channel_prefix_matrix: torch.Tensor + send_head: torch.Tensor + + +@dataclass +class RowMajorInternodeCombineContext: + """Combine context for row-major internode dispatch output.""" + + recv_topk_weights: Optional[torch.Tensor] + src_meta: torch.Tensor + is_token_in_rank: torch.Tensor + recv_rdma_channel_prefix_matrix: torch.Tensor + recv_rdma_rank_prefix_sum: torch.Tensor + recv_gbl_channel_prefix_matrix: torch.Tensor + send_rdma_head: torch.Tensor + send_nvl_head: torch.Tensor + + +CombineContext = Union[ExpertMajorCombineContext, RowMajorIntranodeCombineContext, RowMajorInternodeCombineContext] + + +# Opaque dispatch handles returned by dispatch() and consumed by combine(). + + +@dataclass +class DispatchHandle: + """Base opaque dispatch metadata consumed by :meth:`MoECommunicator.combine`.""" + + output_info: DispatchOutputInfo + + +@dataclass +class ExpertMajorDispatchHandle(DispatchHandle): + combine_context: ExpertMajorCombineContext + + +@dataclass +class RowMajorIntranodeDispatchHandle(DispatchHandle): + combine_context: RowMajorIntranodeCombineContext + + +@dataclass +class RowMajorInternodeDispatchHandle(DispatchHandle): + combine_context: RowMajorInternodeCombineContext + + +# Optional async/overlap configuration. + + +@dataclass +class OperationOverlapConfig: + """Operation-level communication overlap controls.""" + + stream: Optional[torch.cuda.Stream] = None + wait_event: Optional[torch.cuda.Event] = None + num_comm_sms: Optional[int] = None + + +@dataclass +class BlockOverlapConfig: + """Block-level MLP/combine overlap controls.""" + + block_size_m: int + ready_signal: torch.Tensor + ready_value: int = 1 + stream: Optional[torch.cuda.Stream] = None + wait_event: Optional[torch.cuda.Event] = None + num_comm_sms: Optional[int] = None + + +@dataclass +class CommOverlapConfig: + """Mutually exclusive operation-level or block-level overlap configuration.""" + + operation: Optional[OperationOverlapConfig] = None + block: Optional[BlockOverlapConfig] = None + + def __post_init__(self) -> None: + if (self.operation is None) == (self.block is None): + raise ValueError("exactly one of operation or block overlap config must be set") + + @property + def level(self) -> str: + return "block" if self.block is not None else "op" diff --git a/python/mscclpp/ep/utils.py b/python/mscclpp/ep/utils.py new file mode 100644 index 00000000..466bc594 --- /dev/null +++ b/python/mscclpp/ep/utils.py @@ -0,0 +1,148 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Internal helpers shared by the expert-parallel Python frontend.""" + +from __future__ import annotations + +import pickle +from typing import Any, List, Optional, Tuple, Union + +import numpy as np +import torch + + +def send_bytes(comm: Any, payload: bytes, peer: int, tag: int) -> None: + comm.send(np.frombuffer(payload, dtype=np.uint8), peer, tag) + + +def recv_bytes(comm: Any, size: int, peer: int, tag: int) -> bytes: + payload = np.empty(size, dtype=np.uint8) + comm.recv(payload, peer, tag) + return payload.tobytes() + + +def all_gather_object(comm: Any, obj: Any, tag_base: int) -> List[Any]: + payload = pickle.dumps(obj) + rank = comm.my_rank + group_size = comm.nranks + + local_size = np.array([len(payload)], dtype=np.int64) + sizes = np.empty(group_size, dtype=np.int64) + if rank == 0: + sizes[0] = local_size[0] + for peer in range(1, group_size): + comm.recv(sizes[peer : peer + 1], peer, tag_base) + for peer in range(1, group_size): + comm.send(sizes, peer, tag_base + 1) + else: + comm.send(local_size, 0, tag_base) + comm.recv(sizes, 0, tag_base + 1) + + offsets = np.concatenate(([0], np.cumsum(sizes, dtype=np.int64))) + total_size = int(offsets[-1]) + gathered = np.empty(total_size, dtype=np.uint8) + start = int(offsets[rank]) + end = int(offsets[rank + 1]) + if rank == 0: + gathered[start:end] = np.frombuffer(payload, dtype=np.uint8) + for peer in range(1, group_size): + peer_start = int(offsets[peer]) + peer_end = int(offsets[peer + 1]) + comm.recv(gathered[peer_start:peer_end], peer, tag_base + 2) + for peer in range(1, group_size): + comm.send(gathered, peer, tag_base + 3) + else: + send_bytes(comm, payload, 0, tag_base + 2) + comm.recv(gathered, 0, tag_base + 3) + + return [pickle.loads(gathered[int(offsets[i]) : int(offsets[i + 1])].tobytes()) for i in range(group_size)] + + +def broadcast_object(comm: Any, obj: Any, root: int, tag_base: int) -> Any: + rank = comm.my_rank + group_size = comm.nranks + if rank == root: + payload = pickle.dumps(obj) + payload_size = np.array([len(payload)], dtype=np.int64) + for peer in range(group_size): + if peer == root: + continue + comm.send(payload_size, peer, tag_base) + for peer in range(group_size): + if peer == root: + continue + send_bytes(comm, payload, peer, tag_base + 1) + return obj + + payload_size = np.empty(1, dtype=np.int64) + comm.recv(payload_size, root, tag_base) + return pickle.loads(recv_bytes(comm, int(payload_size[0]), root, tag_base + 1)) + + +def ptr(tensor: Optional[torch.Tensor]) -> int: + """``tensor.data_ptr()`` for a tensor, or 0 (== nullptr) for ``None``.""" + return 0 if tensor is None else tensor.data_ptr() + + +def current_stream_ptr() -> int: + """Raw pointer of the current CUDA stream (matches the C++ ``cudaStream_t``).""" + return torch.cuda.current_stream().cuda_stream + + +def cuda_stream_ptr(stream: Optional[torch.cuda.Stream]) -> int: + return (stream if stream is not None else torch.cuda.current_stream()).cuda_stream + + +class DevicePointerArray: + """Minimal ``__cuda_array_interface__`` holder for a non-owning device pointer.""" + + def __init__(self, ptr: int, shape: Tuple[int, ...], typestr: str, owner: Any) -> None: + self._owner = owner + self.__cuda_array_interface__ = { + "data": (ptr, False), + "shape": shape, + "typestr": typestr, + "version": 3, + "strides": None, + } + + +def bf16_view(ptr: int, num_tokens: int, hidden: int, owner: Any) -> torch.Tensor: + """View a raw device pointer as a ``[num_tokens, hidden]`` bfloat16 tensor.""" + u16 = torch.as_tensor(DevicePointerArray(ptr, (num_tokens, hidden), " bool: + fp8_dtype = getattr(torch, "float8_e4m3fn", None) + return fp8_dtype is not None and tensor.dtype == fp8_dtype + + +def exclusive_cumsum(counts: Union[torch.Tensor, List[int]]) -> torch.Tensor: + if isinstance(counts, torch.Tensor): + flat = counts.to(torch.int64).flatten() + zero = torch.zeros(1, dtype=torch.int64, device=flat.device) + return torch.cat([zero, torch.cumsum(flat, dim=0)]) + offsets = [0] + for count in counts: + offsets.append(offsets[-1] + int(count)) + return torch.tensor(offsets, dtype=torch.int64) + + +def resolve_expert_placement( + *, + num_experts: int, + world_size: int, + rank: int, + num_local_experts: Optional[int], + local_expert_start: Optional[int], +) -> Tuple[int, int]: + if num_local_experts is None: + if num_experts % world_size != 0: + raise ValueError("num_experts must be divisible by world_size for even contiguous placement") + num_local_experts = num_experts // world_size + if num_local_experts * world_size != num_experts: + raise NotImplementedError("only even contiguous expert placement is currently supported") + if local_expert_start is None: + local_expert_start = rank * num_local_experts + return num_local_experts, local_expert_start diff --git a/python/mscclpp/ext/__init__.py b/python/mscclpp/ext/__init__.py index 4c8aef5b..08a96ecd 100644 --- a/python/mscclpp/ext/__init__.py +++ b/python/mscclpp/ext/__init__.py @@ -2,9 +2,3 @@ # Licensed under the MIT license. from .algorithm_collection_builder import * - -try: - from . import ep # noqa: F401 -except ImportError: - # EP extension not built; leave `mscclpp.ext.ep` undefined. - pass diff --git a/python/mscclpp/ext/ep/__init__.py b/python/mscclpp/ext/ep/__init__.py deleted file mode 100644 index 04dca708..00000000 --- a/python/mscclpp/ext/ep/__init__.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -"""MSCCL++ Expert-Parallel (MoE dispatch/combine) extension. - -See ``src/ext/ep/README.md`` for migration status and -``python/mscclpp/ext/ep/README.md`` for the high-level API design. - -``MoECommunicator`` is the high-level public API. ``mode=MoEMode.LOW_LATENCY`` -runs on the ``MoERuntime`` LL backend; ``mode=MoEMode.HIGH_THROUGHPUT`` runs on -the DeepEP-style :class:`Buffer` HT backend (GB200 TMA direct-gather combine + -all-sender dispatch). -""" - -from .buffer import Buffer, Config, ExpertParallelRuntime # noqa: F401 -from .communicator import ( # noqa: F401 - CommOverlapConfig, - DispatchHandle, - DispatchLayout, - DispatchOutput, - MoECommunicator, - MoECommunicatorConfig, - MoEMode, - QuantScales, -) - -__all__ = [ - "ExpertParallelRuntime", - "Buffer", - "Config", - "CommOverlapConfig", - "DispatchHandle", - "DispatchLayout", - "DispatchOutput", - "MoECommunicator", - "MoECommunicatorConfig", - "MoEMode", - "QuantScales", -] diff --git a/python/mscclpp/ext/ep/communicator.py b/python/mscclpp/ext/ep/communicator.py deleted file mode 100644 index c530a56b..00000000 --- a/python/mscclpp/ext/ep/communicator.py +++ /dev/null @@ -1,895 +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. - -``MoECommunicator`` is the high-level API. The backend is selected by -:attr:`MoECommunicatorConfig.mode` (a :class:`MoEMode` enum): - -* ``MoEMode.LOW_LATENCY`` — decode path. Wraps the C++ ``MoERuntime`` (RDMA + - CUDA-IPC PortChannel). Output layout ``DispatchLayout.EXPERT_MAJOR``; the - caller pre-allocates the recv buffer. -* ``MoEMode.HIGH_THROUGHPUT`` — prefill path. Wraps the DeepEP-style - :class:`mscclpp.ext.ep.Buffer` (NVLink intranode + RDMA internode HT kernels, - with the GB200 TMA direct-gather combine + all-sender dispatch). Output layout - ``DispatchLayout.FLAT`` grouped by local expert id; intranode vs internode is - selected internally from the RDMA buffer-size hint. - -The two backends are independent C++ runtimes. LL takes an ``mscclpp.CommGroup`` -via ``comm=``; HT takes a ``torch.distributed`` process group via ``group=``. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple, Union - -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 " - "or install with `pip install .[ep]`." - ) from exc - -from .buffer import Config, ExpertParallelRuntime - -DispatchLayout = _cpp.DispatchLayout -MoEMode = _cpp.MoEMode - - -@dataclass -class MoECommunicatorConfig: - """Configuration for the high-level MoE dispatch/combine API.""" - - # Communication. ``comm`` (mscclpp.CommGroup) drives the LL backend; ``group`` - # (torch.distributed ProcessGroup) drives the HT backend. - comm: Optional[Any] = None - group: Optional[dist.ProcessGroup] = None - device: Optional[Union[torch.device, int]] = None - - # Expert topology - num_experts: int = 0 - num_local_experts: Optional[int] = None - local_expert_start: Optional[int] = None - - # 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 - - # Quantization defaults - input_dtype: Optional[torch.dtype] = None - quant_format: Optional[str] = None - - # Transport / launch tuning - num_rdma_qps_per_rank: int = 12 - num_sms: int = 20 - enable_overlap: bool = False - - # HT-only buffer/launch tuning (advanced) - expert_alignment: int = 1 - nvl_chunked_send: int = 8 - nvl_chunked_recv: int = 256 - rdma_chunked_send: int = 16 - rdma_chunked_recv: int = 128 - - -@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: Union[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`. - - LL keeps ``src_info`` / ``layout_range`` (EXPERT_MAJOR). HT keeps the - transport-tagged ``combine_meta`` bundle (FLAT); :meth:`combine` is the only - reader. - """ - - topk_ids: torch.Tensor - weights: torch.Tensor - num_experts: int - num_tokens: int - hidden_size: int - num_local_experts: int - local_expert_start: int - layout: DispatchLayout - output_scales: Optional[QuantScales] = None - # --- LL backend metadata --- - src_info: Optional[torch.Tensor] = None - layout_range: Optional[torch.Tensor] = None - num_max_dispatch_tokens_per_rank: int = 0 - # --- HT backend metadata --- - is_internode: bool = False - combine_meta: Dict[str, Any] = field(default_factory=dict) - - -@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 low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" - - num_sms: int = 20 - - def __init__( - self, - comm: Any, - 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("_MoERuntime supports only MoEMode.LOW_LATENCY") - - 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") - - 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 _CompletionRequest: - """Request handle returned by the HT ``*_async`` methods.""" - - def __init__(self, event, result): - self._event = event - self._result = result - - def wait(self): - if self._event is not None: - try: - self._event.current_stream_wait() - except AttributeError: - torch.cuda.current_stream().synchronize() - return self._result - - -class MoECommunicator: - """High-level MoE communicator for dispatch/combine. - - ``mode=MoEMode.LOW_LATENCY`` selects the LL backend (EXPERT_MAJOR); - ``mode=MoEMode.HIGH_THROUGHPUT`` selects the HT backend (FLAT). - """ - - 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: - config = MoECommunicatorConfig(**kwargs) - - if config.device is not None: - torch.cuda.set_device(config.device) - - if not isinstance(config.mode, MoEMode): - raise TypeError("MoECommunicatorConfig.mode must be a MoEMode") - self.mode = config.mode - self.output_layout = _resolve_output_layout(config.output_layout, self.mode) - - # ---- shared shape / capacity ---- - 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") - - self.num_sms = config.num_sms - self.enable_overlap = config.enable_overlap - - if self.mode == MoEMode.LOW_LATENCY: - self._init_ll(config) - else: - self._init_ht(config) - - # ------------------------------------------------------------------ - # Backend construction - # ------------------------------------------------------------------ - - def _resolve_placement(self) -> None: - if self.num_local_experts is None: - if self.num_experts % self.world_size != 0: - raise ValueError("num_experts must be divisible by world_size for even contiguous placement") - 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") - if self.local_expert_start is None: - self.local_expert_start = self.rank * self.num_local_experts - - def _init_ll(self, config: MoECommunicatorConfig) -> None: - from mscclpp._core import CommGroup # local import: only LL needs it - - comm = config.comm - if comm is None: - if config.group is None: - raise ValueError("mode=LOW_LATENCY requires an mscclpp.CommGroup via comm= (or a torch group=)") - comm = CommGroup(torch_group=config.group) - self.comm = comm - self.rank = comm.my_rank - self.world_size = comm.nranks - self.local_rank = torch.cuda.current_device() - self.device = torch.device("cuda", self.local_rank) - - if self.output_layout != DispatchLayout.EXPERT_MAJOR: - raise NotImplementedError("low-latency mode currently supports only DispatchLayout.EXPERT_MAJOR") - 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 - self.local_expert_start = config.local_expert_start - self._resolve_placement() - - 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_rdma_bytes = _get_low_latency_rdma_size_hint( - self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts - ) - 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=0, - num_rdma_bytes=num_rdma_bytes, - mode=self.mode, - num_qps_per_rank=config.num_rdma_qps_per_rank, - ) - # Report internode only for genuine multi-node (multi-NVLink-domain) jobs - # instead of hard-coding True: LL always uses the RDMA transport, but a - # single-node LL job is not internode topology-wise. num_rdma_ranks > 1 - # iff world_size spans more than one local NVLink domain. - self._is_internode = self._runtime.get_num_rdma_ranks() > 1 - - def _init_ht(self, config: MoECommunicatorConfig) -> None: - group = config.group - if group is None: - raise ValueError("mode=HIGH_THROUGHPUT requires a torch.distributed ProcessGroup via group=") - self.group = group - self.rank = group.rank() - self.world_size = group.size() - self.local_rank = torch.cuda.current_device() - self.device = torch.device("cuda", self.local_rank) - - if self.output_layout != DispatchLayout.FLAT: - raise NotImplementedError("HT mode currently supports only DispatchLayout.FLAT") - - self.num_local_experts = config.num_local_experts - self.local_expert_start = config.local_expert_start - self._resolve_placement() - - if config.input_dtype not in (None, torch.bfloat16): - raise NotImplementedError("HT dispatch currently supports BF16 input only") - if config.quant_format is not None: - raise NotImplementedError("HT quantized dispatch (scales) is not implemented yet") - - self.expert_alignment = config.expert_alignment - - # Config(num_sms, nvl_send, nvl_recv, rdma_send, rdma_recv). The C++ size - # hints return 0 RDMA bytes when world_size <= NUM_MAX_NVL_PEERS, which is - # exactly the intranode/internode boundary — so derive the transport from - # the hint instead of hardcoding the NVL peer count. - self._cfg = Config( - self.num_sms, - config.nvl_chunked_send, - config.nvl_chunked_recv, - config.rdma_chunked_send, - config.rdma_chunked_recv, - ) - hidden_bytes = self.hidden_size * torch.tensor([], dtype=torch.bfloat16).element_size() - num_nvl_bytes = self._cfg.get_nvl_buffer_size_hint(hidden_bytes, self.world_size) - num_rdma_bytes = self._cfg.get_rdma_buffer_size_hint(hidden_bytes, self.world_size) - self._is_internode = num_rdma_bytes > 0 - - self._buffer = ExpertParallelRuntime( - group, - num_nvl_bytes=num_nvl_bytes, - num_rdma_bytes=num_rdma_bytes, - low_latency_mode=False, - num_qps_per_rank=config.num_rdma_qps_per_rank, - ) - - # ------------------------------------------------------------------ - # Introspection - # ------------------------------------------------------------------ - - def is_available(self) -> bool: - return self._runtime.is_available() if self.mode == MoEMode.LOW_LATENCY else self._buffer.is_available() - - def is_internode_available(self) -> bool: - if self.mode == MoEMode.LOW_LATENCY: - return self._runtime.is_internode_available() - return self._buffer.is_internode_available() - - def is_internode(self) -> bool: - return self._is_internode - - # ------------------------------------------------------------------ - # Dispatch - # ------------------------------------------------------------------ - - def dispatch( - self, - input: torch.Tensor, - topk_ids: torch.Tensor, - weights: Optional[torch.Tensor] = None, - scales: Optional[QuantScales] = None, - *, - output_buffer: Optional[torch.Tensor] = None, - stream: Optional[torch.cuda.Stream] = None, - previous_handle: Optional[DispatchHandle] = None, - ) -> Tuple[DispatchOutput, DispatchHandle]: - if self.mode == MoEMode.LOW_LATENCY: - return self._dispatch_ll(input, topk_ids, weights, scales, output_buffer, stream) - # HT honors a caller-provided stream by running its kernels on it (the HT - # runtime reads torch.cuda.current_stream()); routing the whole call - # through a stream context keeps the two-phase notify/dispatch ordered. - if stream is not None: - with torch.cuda.stream(stream): - return self._dispatch_ht(input, topk_ids, weights, scales, previous_handle) - return self._dispatch_ht(input, topk_ids, weights, scales, previous_handle) - - def _dispatch_ll(self, input, topk_ids, weights, scales, output_buffer, stream): - 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) - - out_buf, packed_scales, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) - self._runtime._cpp_runtime.dispatch( - input.data_ptr(), - topk_ids.data_ptr(), - out_buf.data_ptr(), - 0 if packed_scales is None else packed_scales.data_ptr(), - src_info.data_ptr(), - layout_range.data_ptr(), - count.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=out_buf, - scales=output_scales, - num_tokens_per_expert=count, - expert_offsets=None, - layout=self.output_layout, - ) - handle = DispatchHandle( - topk_ids=topk_ids, - weights=weights, - 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, - src_info=src_info, - layout_range=layout_range, - num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank, - ) - return dispatch_out, handle - - def _dispatch_ht(self, input, topk_ids, weights, scales, previous_handle): - self._validate_dispatch_inputs_ht(input, topk_ids, weights, scales) - if weights is None: - weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) - - cache = getattr(previous_handle, "_dispatch_cache", None) if previous_handle is not None else None - if cache is not None: - num_tokens_per_rank = cache["num_tokens_per_rank"] - num_tokens_per_rdma_rank = cache["num_tokens_per_rdma_rank"] - num_tokens_per_expert = cache["num_tokens_per_expert"] - is_token_in_rank = cache["is_token_in_rank"] - else: - ( - num_tokens_per_rank, - num_tokens_per_rdma_rank, - num_tokens_per_expert, - is_token_in_rank, - ) = self._buffer.get_dispatch_layout(topk_ids, self.num_experts) - - if self._is_internode: - ( - recv_x, - _recv_x_scales, - _recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - _rdma_channel_prefix_matrix, - _gbl_channel_prefix_matrix, - recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix, - _recv_gbl_rank_prefix_sum, - recv_src_meta, - send_rdma_head, - send_nvl_head, - ) = self._buffer.internode_dispatch( - input, - None, - topk_ids, - weights, - num_tokens_per_rank, - num_tokens_per_rdma_rank, - is_token_in_rank, - num_tokens_per_expert, - 0, - 0, - None, - None, - None, - None, - self.expert_alignment, - self._cfg, - ) - combine_meta = { - "recv_topk_weights": recv_topk_weights, - "src_meta": recv_src_meta, - "is_token_in_rank": is_token_in_rank, - "recv_rdma_channel_prefix_matrix": recv_rdma_channel_prefix_matrix, - "recv_rdma_rank_prefix_sum": recv_rdma_rank_prefix_sum, - "recv_gbl_channel_prefix_matrix": recv_gbl_channel_prefix_matrix, - "send_rdma_head": send_rdma_head, - "send_nvl_head": send_nvl_head, - } - dispatch_cache = ( - cache - if cache is not None - else { - "num_tokens_per_rank": num_tokens_per_rank, - "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, - "num_tokens_per_expert": num_tokens_per_expert, - "is_token_in_rank": is_token_in_rank, - } - ) - elif cache is not None: - ( - recv_x, - _recv_x_scales, - _recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rank_prefix_matrix, - _channel_prefix_matrix, - recv_channel_prefix_matrix, - recv_src_idx, - send_head, - ) = self._buffer.intranode_dispatch( - input, - None, - None, - None, - None, - is_token_in_rank, - None, - cache["num_recv_tokens"], - cache["rank_prefix_matrix"], - cache["channel_prefix_matrix"], - self.expert_alignment, - self._cfg, - ) - combine_meta = { - "recv_topk_weights": recv_topk_weights, - "src_idx": recv_src_idx, - "rank_prefix_matrix": rank_prefix_matrix, - "recv_channel_prefix_matrix": recv_channel_prefix_matrix, - "send_head": send_head, - } - dispatch_cache = cache - else: - ( - recv_x, - _recv_x_scales, - _recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rank_prefix_matrix, - channel_prefix_matrix, - recv_channel_prefix_matrix, - recv_src_idx, - send_head, - ) = self._buffer.intranode_dispatch( - input, - None, - topk_ids, - weights, - num_tokens_per_rank, - is_token_in_rank, - num_tokens_per_expert, - 0, - None, - None, - self.expert_alignment, - self._cfg, - ) - combine_meta = { - "recv_topk_weights": recv_topk_weights, - "src_idx": recv_src_idx, - "rank_prefix_matrix": rank_prefix_matrix, - "recv_channel_prefix_matrix": recv_channel_prefix_matrix, - "send_head": send_head, - } - dispatch_cache = { - "num_tokens_per_rank": num_tokens_per_rank, - "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, - "num_tokens_per_expert": num_tokens_per_expert, - "is_token_in_rank": is_token_in_rank, - "rank_prefix_matrix": rank_prefix_matrix, - "channel_prefix_matrix": channel_prefix_matrix, - "num_recv_tokens": int(recv_x.size(0)), - } - - dispatch_out = DispatchOutput( - tokens=recv_x, - scales=None, - num_tokens_per_expert=num_recv_tokens_per_expert_list, - expert_offsets=_exclusive_cumsum(num_recv_tokens_per_expert_list), - layout=DispatchLayout.FLAT, - ) - handle = DispatchHandle( - topk_ids=topk_ids, - weights=weights, - num_experts=self.num_experts, - num_tokens=int(input.size(0)), - hidden_size=self.hidden_size, - num_local_experts=self.num_local_experts, - local_expert_start=self.local_expert_start, - layout=DispatchLayout.FLAT, - output_scales=None, - is_internode=self._is_internode, - combine_meta=combine_meta, - ) - # The torch-free HT runtime orders its work on the caller's CUDA stream - # (no separate event handle), so there is nothing to attach here. - handle._event = None # type: ignore[attr-defined] - handle._dispatch_cache = dispatch_cache # type: ignore[attr-defined] - return dispatch_out, handle - - # ------------------------------------------------------------------ - # Combine - # ------------------------------------------------------------------ - - def combine( - self, - expert_output: torch.Tensor, - handle: DispatchHandle, - *, - out: Optional[torch.Tensor] = None, - stream: Optional[torch.cuda.Stream] = None, - ) -> torch.Tensor: - if self.mode == MoEMode.LOW_LATENCY: - return self._combine_ll(expert_output, handle, out, stream) - # HT honors a caller-provided stream by running its kernels on it (see dispatch()). - if stream is not None: - with torch.cuda.stream(stream): - return self._combine_ht(expert_output, handle, out) - return self._combine_ht(expert_output, handle, out) - - def _combine_ll(self, expert_output, handle, out, stream): - 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 _combine_ht(self, expert_output, handle, out): - self._validate_combine_inputs_ht(expert_output, handle) - m = handle.combine_meta - if handle.is_internode: - combined_x, _combined_w = self._buffer.internode_combine( - expert_output, - m["recv_topk_weights"], - m["src_meta"], - m["is_token_in_rank"], - m["recv_rdma_channel_prefix_matrix"], - m["recv_rdma_rank_prefix_sum"], - m["recv_gbl_channel_prefix_matrix"], - m["send_rdma_head"], - m["send_nvl_head"], - self._cfg, - ) - else: - combined_x, _combined_w = self._buffer.intranode_combine( - expert_output, - m["recv_topk_weights"], - m["src_idx"], - m["rank_prefix_matrix"], - m["recv_channel_prefix_matrix"], - m["send_head"], - self._cfg, - ) - if out is not None: - out.copy_(combined_x) - return out - return combined_x - - # ------------------------------------------------------------------ - # Optional async / overlap APIs (HT only) - # ------------------------------------------------------------------ - - 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) - - # ------------------------------------------------------------------ - # LL output tensors + validation - # ------------------------------------------------------------------ - - def _get_dispatch_output_tensors(self, output_buffer: 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) - return ( - output_buffer, - self._dispatch_scales, - self._dispatch_src_info, - self._dispatch_layout_range, - self._dispatch_count, - ) - - def _validate_dispatch_inputs(self, input, topk_ids, weights, scales, output_buffer) -> 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, handle, out) -> 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}") - - # ------------------------------------------------------------------ - # HT validation - # ------------------------------------------------------------------ - - def _validate_dispatch_inputs_ht(self, input, topk_ids, weights, scales) -> None: - if scales is not None and (scales.local is not None or scales.global_scale is not None): - raise NotImplementedError("HT 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] tensor") - if input.device.type != "cuda" or input.dtype != torch.bfloat16: - raise ValueError("HT dispatch input must be a CUDA BF16 tensor") - if input.size(1) != self.hidden_size: - raise ValueError(f"input hidden size {input.size(1)} != 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 be [input.size(0), 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") - - def _validate_combine_inputs_ht(self, expert_output, handle) -> None: - if not isinstance(handle, DispatchHandle): - raise TypeError("handle must be a DispatchHandle returned by dispatch") - if expert_output.dim() != 2 or not expert_output.is_contiguous(): - raise ValueError("expert_output must be a contiguous [total_recv_tokens, hidden] tensor") - if expert_output.size(1) != self.hidden_size: - raise ValueError(f"expert_output hidden size {expert_output.size(1)} != configured {self.hidden_size}") - if handle.is_internode != self._is_internode: - raise ValueError("handle transport does not match this communicator") - - -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 _exclusive_cumsum(counts: Union[torch.Tensor, List[int]]) -> torch.Tensor: - if isinstance(counts, torch.Tensor): - flat = counts.to(torch.int64).flatten() - zero = torch.zeros(1, dtype=torch.int64, device=flat.device) - return torch.cat([zero, torch.cumsum(flat, dim=0)]) - offsets = [0] - for c in counts: - offsets.append(offsets[-1] + int(c)) - return torch.tensor(offsets, dtype=torch.int64) - - -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/src/ext/ep/README.md b/src/ext/ep/README.md index d80cf746..c6ac8e4a 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -18,7 +18,7 @@ MSCCL++. The module builds two active backends: | HT dispatch/combine | ✅ active DeepEP-style backend with intranode and internode paths | | HT GB200 direct/flat fast paths | ✅ runtime-gated by `MSCCLPP_EP_DIRECT`, `MSCCLPP_EP_INTRA_DIRECT`, and `MSCCLPP_EP_FLAT` | | GB200 LL NVLS multimem fast path | ✅ runtime-gated by `mscclpp::isNvlsSupported()` | -| Python frontend `mscclpp.ext.ep` | ✅ `MoECommunicator` selects LL or HT by `MoEMode` | +| Python frontend `mscclpp.ep` | ✅ `MoECommunicator` selects LL or HT by `MoEMode` | 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, @@ -102,7 +102,7 @@ This produces `mscclpp_ep_cpp.so` — a nanobind extension module. The Python frontend picks it up automatically: ```python -from mscclpp.ext import ep +import mscclpp.ep as ep moe_comm = ep.MoECommunicator(...) ``` @@ -324,11 +324,11 @@ src/ext/ep/ ├── utils.cuh — device inline helpers └── low_latency.cu — LL dispatch/combine (RDMA + IPC paths) -python/mscclpp/ext/ep/ +python/mscclpp/ep/ ├── __init__.py — reexports the public MoECommunicator API └── communicator.py — torch.Tensor frontend over raw-pointer runtime calls -test/python/ext/ep/ +test/python/ep/ ├── test_intranode_multirank.py — intranode HT dispatch+combine ├── test_internode_multirank.py — internode HT dispatch+combine └── test_low_latency_multirank.py — LL dispatch+combine @@ -364,7 +364,7 @@ pip install scikit-build-core nanobind setuptools_scm Then build the EP extension (see [Build](#build)) — `pip install .` from the repo root installs `mscclpp_ep_cpp.so` into the active env so the -test scripts can `from mscclpp.ext import ep`. +test scripts can `import mscclpp.ep as ep`. Intranode (single node, 8 GPUs) — HT: @@ -373,7 +373,7 @@ MSCCLPP_EP_BENCH=1 \ MSCCLPP_EP_BENCH_TOKENS=4096 MSCCLPP_EP_BENCH_HIDDEN=7168 \ MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ torchrun --nnodes=1 --nproc_per_node=8 \ - test/python/ext/ep/test_intranode_multirank.py + test/python/ep/test_intranode_multirank.py ``` Intranode LL (single node, 8 GPUs): @@ -383,7 +383,7 @@ MSCCLPP_EP_BENCH=1 \ MSCCLPP_EP_BENCH_TOKENS=128 MSCCLPP_EP_BENCH_HIDDEN=7168 \ MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ torchrun --nnodes=1 --nproc_per_node=8 \ - test/python/ext/ep/test_low_latency_multirank.py + test/python/ep/test_low_latency_multirank.py ``` Internode HT (2 nodes × 8 GPUs), torchrun: @@ -395,7 +395,7 @@ MSCCLPP_EP_BENCH_TOKENS=4096 MSCCLPP_EP_BENCH_HIDDEN=7168 \ MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 \ --master_addr= --master_port=29600 \ - test/python/ext/ep/test_internode_multirank.py + test/python/ep/test_internode_multirank.py # node 1 (worker) MSCCLPP_EP_BENCH=1 \ @@ -403,7 +403,7 @@ MSCCLPP_EP_BENCH_TOKENS=4096 MSCCLPP_EP_BENCH_HIDDEN=7168 \ MSCCLPP_EP_BENCH_EXPERTS=256 MSCCLPP_EP_BENCH_TOPK=8 \ torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 \ --master_addr= --master_port=29600 \ - test/python/ext/ep/test_internode_multirank.py + test/python/ep/test_internode_multirank.py ``` If the bootstrap NIC is mis-detected (e.g. multi-homed hosts), pin @@ -427,7 +427,7 @@ mpirun -np 16 --allow-run-as-root --hostfile \ bash -c 'export RANK=$OMPI_COMM_WORLD_RANK \ WORLD_SIZE=$OMPI_COMM_WORLD_SIZE \ LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK; \ - exec python3 test/python/ext/ep/test_internode_multirank.py' + exec python3 test/python/ep/test_internode_multirank.py' ``` Internode LL via mpirun — same launch wrapper, swap the test script: @@ -442,7 +442,7 @@ mpirun -np 16 --allow-run-as-root --hostfile \ bash -c 'export RANK=$OMPI_COMM_WORLD_RANK \ WORLD_SIZE=$OMPI_COMM_WORLD_SIZE \ LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK; \ - exec python3 test/python/ext/ep/test_low_latency_multirank.py' + exec python3 test/python/ep/test_low_latency_multirank.py' ``` Add `-x NCCL_SOCKET_IFNAME= -x MSCCLPP_SOCKET_IFNAME= diff --git a/src/ext/ep/ht_runtime.cc b/src/ext/ep/ht_runtime.cc index afd01337..e7e26994 100644 --- a/src/ext/ep/ht_runtime.cc +++ b/src/ext/ep/ht_runtime.cc @@ -1273,7 +1273,7 @@ void MoEHighThroughputRuntime::intranodeCombine(void* combinedX, float* combined // Internode (NVLink + RDMA) high-throughput path. Ported from DeepEP // `csrc/deep_ep.cpp`; the kernels it drives are in // `src/ext/ep/kernels/internode.cu`. Validated end-to-end on 2 x H100 x 8 -// via `test/python/ext/ep/test_internode_multirank.py`. De-torched the same +// via `test/python/ep/test_internode_multirank.py`. De-torched the same // way as the intranode path: tensor params became raw pointers + size ints, // output tensors became caller pointers, the EventHandle / async / record_stream // machinery became comm_stream stream_wait brackets, and the original single diff --git a/test/python/ext/ep/test_internode_multirank.py b/test/python/ep/test_internode_multirank.py similarity index 75% rename from test/python/ext/ep/test_internode_multirank.py rename to test/python/ep/test_internode_multirank.py index bda20367..3569b392 100644 --- a/test/python/ext/ep/test_internode_multirank.py +++ b/test/python/ep/test_internode_multirank.py @@ -8,13 +8,13 @@ Launch on each node with (example: 2 nodes x 8 GPUs = 16 ranks): MASTER_ADDR= MASTER_PORT=29600 NODE_RANK=0 \ torchrun --nnodes=2 --nproc_per_node=8 \ --rdzv-backend=c10d --rdzv-endpoint=:29600 \ - test/python/ext/ep/test_internode_multirank.py + test/python/ep/test_internode_multirank.py # on worker (NODE_RANK=1): MASTER_ADDR= MASTER_PORT=29600 NODE_RANK=1 \ torchrun --nnodes=2 --nproc_per_node=8 \ --rdzv-backend=c10d --rdzv-endpoint=:29600 \ - test/python/ext/ep/test_internode_multirank.py + test/python/ep/test_internode_multirank.py Round-trip dispatch + combine using internode HT kernels across nodes. @@ -86,7 +86,7 @@ def inplace_unique(x: torch.Tensor, num_slots: int): def main(): rank, num_ranks, local_rank, group = init_dist() from mscclpp import CommGroup - from mscclpp.ext import ep + import mscclpp.ep as ep ep_group = CommGroup(torch_group=group) @@ -142,138 +142,59 @@ def main(): x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - # 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")), - int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")), - int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")), - int(os.environ.get("MSCCLPP_EP_RDMA_SEND", "16")), - int(os.environ.get("MSCCLPP_EP_RDMA_RECV", "128")), + moe = ep.MoECommunicator( + comm=ep_group, + num_experts=num_experts, + hidden_size=hidden, + topk=num_topk, + max_tokens_per_rank=num_tokens, + mode=ep.MoEMode.HIGH_THROUGHPUT, + num_sms=int(os.environ.get("MSCCLPP_EP_NSM", "152")), + nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")), + nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")), + rdma_chunked_send=int(os.environ.get("MSCCLPP_EP_RDMA_SEND", "16")), + rdma_chunked_recv=int(os.environ.get("MSCCLPP_EP_RDMA_RECV", "128")), ) - _bench_on = os.environ.get("MSCCLPP_EP_BENCH", "0") == "1" - _buf_hidden = max(hidden, int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "0"))) if _bench_on else hidden - num_nvl_bytes = cfg.get_nvl_buffer_size_hint(_buf_hidden * x.element_size(), num_ranks) - num_rdma_bytes = cfg.get_rdma_buffer_size_hint(_buf_hidden * x.element_size(), num_ranks) if rank == 0: print( f"[cfg] num_nodes={num_nodes} num_ranks={num_ranks} num_tokens={num_tokens} " - f"hidden={hidden} num_experts={num_experts} num_topk={num_topk} " - f"num_nvl_bytes={num_nvl_bytes} num_rdma_bytes={num_rdma_bytes}", + f"hidden={hidden} num_experts={num_experts} num_topk={num_topk}", flush=True, ) - print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True) - buf = ep.ExpertParallelRuntime( - ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False - ) print( - f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()} " - f"is_internode={buf.is_internode_available()}", + f"[rank {rank}] MoECommunicator created is_available={moe.is_available()} " + f"is_internode={moe.is_internode_available()}", flush=True, ) - assert buf.is_available() and buf.is_internode_available() + assert moe.is_available() and moe.is_internode_available() + assert moe.is_internode(), "expected the communicator to select the internode HT transport" - ref_rank, ref_rdma_rank, ref_exp, ref_in_rank, _ = buf.get_dispatch_layout( - topk_idx, num_experts, None, False, False - ) - assert torch.allclose(ref_rank, num_tokens_per_rank) - assert torch.allclose(ref_rdma_rank, num_tokens_per_rdma_rank) - assert torch.allclose(ref_exp, num_tokens_per_expert) - assert torch.allclose(ref_in_rank, is_token_in_rank) - if rank == 0: - print("[layout] OK", flush=True) - dist.barrier(group=group) - - # internode_dispatch signature (non-cached mode): - # (x, x_scales, topk_idx, topk_weights, - # num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank, num_tokens_per_expert, - # cached_num_recv_tokens=0, cached_num_rdma_recv_tokens=0, - # cached_rdma_channel_prefix_matrix=None, cached_recv_rdma_rank_prefix_sum=None, - # cached_gbl_channel_prefix_matrix=None, cached_recv_gbl_rank_prefix_sum=None, - # expert_alignment, config) - ( - recv_x, - recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rdma_channel_prefix_matrix, - gbl_channel_prefix_matrix, - recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix, - recv_gbl_rank_prefix_sum, - recv_src_meta, - send_rdma_head, - send_nvl_head, - ) = buf.internode_dispatch( + dispatch_out, handle = moe.dispatch( x, - None, topk_idx, topk_weights, - num_tokens_per_rank, - num_tokens_per_rdma_rank, - is_token_in_rank, - num_tokens_per_expert, - 0, - 0, - None, - None, - None, - None, - 1, - cfg, ) + recv_x = dispatch_out.tokens dist.barrier(group=group) - _skip_verify = os.environ.get("MSCCLPP_EP_SKIP_VERIFY", "0") in ("1", "true", "True") - # Validate recv buffer: for each source rank i, the block carries value i. assert recv_x.dim() == 2 and recv_x.size(1) == hidden - start = 0 - for src in range(num_ranks): - end = recv_gbl_rank_prefix_sum[src].item() - block = recv_x[start:end] - if block.numel(): - lo = block.float().amin().item() - hi = block.float().amax().item() - assert _skip_verify or ( - abs(lo - src) < 1e-3 and abs(hi - src) < 1e-3 - ), f"rank{rank}: block from src={src} has range=[{lo}, {hi}], expected {src}" - start = end + local_experts = num_experts // num_ranks + all_expert_counts = torch.empty((num_ranks, num_experts), dtype=num_tokens_per_expert.dtype, device="cuda") + dist.all_gather_into_tensor(all_expert_counts, num_tokens_per_expert, group=group) + expected_counts = all_expert_counts[:, rank * local_experts : (rank + 1) * local_experts].sum(dim=0).cpu().tolist() + assert dispatch_out.layout.num_tokens_per_expert is not None + actual_counts = [int(count) for count in dispatch_out.layout.num_tokens_per_expert] + assert actual_counts == [int(count) for count in expected_counts] if rank == 0: print(f"[dispatch] OK (recv {recv_x.size(0)} tokens)", flush=True) - # XXX: forcing a device+group sync here is currently required for combine - # to see consistent dispatch outputs. Without this both send_nvl_head and - # 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 ExpertParallelRuntime.internode_dispatch. + # Keep the existing dispatch/combine phase guard for internode HT until the + # backend wires a proper stream-dependency hand-off. torch.cuda.synchronize() dist.barrier(group=group) - # internode_combine signature: - # (x, topk_weights, - # src_meta, is_combined_token_in_rank, - # rdma_channel_prefix_matrix, rdma_rank_prefix_sum, gbl_channel_prefix_matrix, - # combined_rdma_head, combined_nvl_head, config) - # NOTE: combine goes in the reverse direction of dispatch, so the prefix - # matrices passed here must be the RECEIVER-side ones returned by dispatch - # (`recv_rdma_channel_prefix_matrix`, `recv_rdma_rank_prefix_sum`, - # `recv_gbl_channel_prefix_matrix`) — not the sender-side ones. - combined_x, combined_topk_weights = buf.internode_combine( - recv_x, - recv_topk_weights, - recv_src_meta, - is_token_in_rank, - recv_rdma_channel_prefix_matrix, - recv_rdma_rank_prefix_sum, - recv_gbl_channel_prefix_matrix, - send_rdma_head, - send_nvl_head, - cfg, - ) + combined_x = moe.combine(recv_x, handle) num_dst = is_token_in_rank.sum(dim=1).to(torch.float32) expected = num_dst * float(rank) @@ -284,7 +205,7 @@ def main(): # bf16 accumulator has 7-bit mantissa; intermediate partial sums can # round at ulp = max_exp * 2**-7. Use a tolerance that scales with magnitude. tol = max(1e-2, max_exp * (1.0 / 64)) - assert _skip_verify or diff <= tol, f"rank{rank}: combine mismatch max diff {diff} > tol {tol} (max_exp={max_exp})" + assert diff <= tol, f"rank{rank}: combine mismatch max diff {diff} > tol {tol} (max_exp={max_exp})" dist.barrier(group=group) if rank == 0: @@ -317,19 +238,6 @@ def main(): print(f"[bench] skip: topk={bench_num_topk} > experts={bench_num_experts}", flush=True) return - # 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): - if rank == 0: - 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 runtime or smaller hidden.", - flush=True, - ) - return - scores_b = torch.randn((bench_tokens, bench_num_experts), device="cuda", dtype=torch.float32).abs() + 1 topk_idx_b = torch.topk(scores_b, bench_num_topk, dim=-1, sorted=False).indices topk_weights_b = torch.ones((bench_tokens, bench_num_topk), dtype=torch.float32, device="cuda") @@ -359,17 +267,14 @@ def main(): is_token_in_rank_b = token_idx_in_rank_b >= 0 x_b = torch.ones((bench_tokens, bench_hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - # Drive the benchmark through the high-level MoECommunicator (the public - # #818 API), mode=HIGH_THROUGHPUT. With world_size > NUM_MAX_NVL_PEERS the - # RDMA size hint is non-zero, so the communicator auto-selects the internode - # transport (internode_dispatch / internode_combine) internally. It owns its - # own ExpertParallelRuntime sized for the bench shape and runs - # get_dispatch_layout internally on the first (uncached) dispatch, recording - # the routing layout on the returned handle; subsequent dispatches reuse it - # via previous_handle, skipping the host-side layout computation. This - # isolates the on-GPU dispatch-kernel cost (NCCL-EP ep_bench convention). + # Drive the benchmark through the public high-level API. The communicator + # auto-selects internode HT when the RDMA size hint is non-zero. The first + # (uncached) dispatch records routing layout on the returned handle; + # subsequent dispatches reuse it via previous_handle, skipping host-side + # layout computation. This isolates the on-GPU dispatch-kernel cost + # (NCCL-EP ep_bench convention). moe = ep.MoECommunicator( - group=group, + comm=ep_group, num_experts=bench_num_experts, hidden_size=bench_hidden, topk=bench_num_topk, diff --git a/test/python/ext/ep/test_intranode_multirank.py b/test/python/ep/test_intranode_multirank.py similarity index 79% rename from test/python/ext/ep/test_intranode_multirank.py rename to test/python/ep/test_intranode_multirank.py index fd87c585..977284df 100644 --- a/test/python/ext/ep/test_intranode_multirank.py +++ b/test/python/ep/test_intranode_multirank.py @@ -3,10 +3,11 @@ """Multi-rank intranode functional validation for mscclpp_ep. Launch with: - torchrun --nproc_per_node= test/python/ext/ep/test_intranode_multirank.py + torchrun --nproc_per_node= test/python/ep/test_intranode_multirank.py -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). +Tests that the high-level ``MoECommunicator`` succeeds across N GPUs on a single +node and that a round-trip dispatch + combine preserves data (sum of top-k +weighted copies). Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass that times dispatch and combine **separately** with CUDA events and @@ -66,7 +67,7 @@ def inplace_unique(x: torch.Tensor, num_slots: int): def main(): rank, num_ranks, local_rank, group = init_dist() from mscclpp import CommGroup - from mscclpp.ext import ep + import mscclpp.ep as ep ep_group = CommGroup(torch_group=group) @@ -107,95 +108,46 @@ 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 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")), - int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")), - int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")), + moe = ep.MoECommunicator( + comm=ep_group, + num_experts=num_experts, + hidden_size=hidden, + topk=num_topk, + max_tokens_per_rank=num_tokens, + mode=ep.MoEMode.HIGH_THROUGHPUT, + num_sms=int(os.environ.get("MSCCLPP_EP_NUM_SMS", "20")), + nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")), + nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")), ) - _bench_on = os.environ.get("MSCCLPP_EP_BENCH", "0") == "1" - _buf_hidden = max(hidden, int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "0"))) if _bench_on else hidden - num_nvl_bytes = cfg.get_nvl_buffer_size_hint(_buf_hidden * x.element_size(), num_ranks) 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} num_nvl_bytes={num_nvl_bytes}", + f"num_experts={num_experts} num_topk={num_topk}", flush=True, ) + print(f"[rank {rank}] MoECommunicator created is_available={moe.is_available()}", flush=True) + assert moe.is_available() - 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 - ref_rank, _, ref_exp, ref_in_rank = buf.get_dispatch_layout(topk_idx, num_experts) - assert torch.allclose(ref_rank, num_tokens_per_rank) - assert torch.allclose(ref_exp, num_tokens_per_expert) - assert torch.allclose(ref_in_rank, is_token_in_rank) - if rank == 0: - print("[layout] OK", flush=True) - - # Dispatch - ( - recv_x, - recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rank_prefix_matrix, - channel_prefix_matrix, - recv_channel_prefix_matrix, - recv_src_idx, - send_head, - ) = buf.intranode_dispatch( + dispatch_out, handle = moe.dispatch( x, - None, topk_idx, topk_weights, - num_tokens_per_rank, - is_token_in_rank, - num_tokens_per_expert, - 0, - None, - None, - 1, - cfg, ) + recv_x = dispatch_out.tokens dist.barrier(group=group) - # Validate received payloads: for each source rank i, the block of tokens - # we received from it should be filled with `i`. assert recv_x.dim() == 2 and recv_x.size(1) == hidden - start = 0 - for src in range(num_ranks): - end = rank_prefix_matrix[src][rank].item() - block = recv_x[start:end] - if block.numel(): - actual = block.float().amin().item() - assert abs(actual - src) < 1e-3, f"rank{rank}: block from src={src} has min={actual}, expected {src}" - assert abs(block.float().amax().item() - src) < 1e-3 - start = end + local_experts = num_experts // num_ranks + all_expert_counts = torch.empty((num_ranks, num_experts), dtype=num_tokens_per_expert.dtype, device="cuda") + dist.all_gather_into_tensor(all_expert_counts, num_tokens_per_expert, group=group) + expected_counts = all_expert_counts[:, rank * local_experts : (rank + 1) * local_experts].sum(dim=0).cpu().tolist() + assert dispatch_out.layout.num_tokens_per_expert is not None + actual_counts = [int(count) for count in dispatch_out.layout.num_tokens_per_expert] + assert actual_counts == [int(count) for count in expected_counts] if rank == 0: print(f"[dispatch] OK (recv {recv_x.size(0)} tokens)", flush=True) - # Combine (scatter-reduce back). Using recv_topk_weights=None path with - # dispatched tokens unchanged => every source rank should receive its - # contribution back, unweighted sum across topk copies. - handle_recv_src_idx = recv_src_idx - handle_rank_prefix_matrix = rank_prefix_matrix - handle_channel_prefix_matrix = recv_channel_prefix_matrix - - combined_x, combined_topk_weights = buf.intranode_combine( - recv_x, - recv_topk_weights, - handle_recv_src_idx, - handle_rank_prefix_matrix, - handle_channel_prefix_matrix, - send_head, - cfg, - ) + combined_x = moe.combine(recv_x, handle) # Expected: we dispatched with x = rank * ones, so every destination r # received the value `rank` for our token. On combine the destinations @@ -241,19 +193,8 @@ def main(): print(f"[bench] skip: topk={bench_num_topk} > experts={bench_num_experts}", flush=True) return - # Rebuild inputs at bench size. Keep same layout recipe as above but at - # 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 runtime or smaller hidden.", - flush=True, - ) - return - + # Rebuild inputs at bench size. The benchmark creates its own communicator + # below so its internal buffers are sized for the benchmark shape. scores_b = torch.randn((bench_tokens, bench_num_experts), device="cuda", dtype=torch.float32).abs() + 1 topk_idx_b = torch.topk(scores_b, bench_num_topk, dim=-1, sorted=False).indices topk_weights_b = torch.ones((bench_tokens, bench_num_topk), dtype=torch.float32, device="cuda") @@ -276,15 +217,13 @@ def main(): is_token_in_rank_b = token_idx_in_rank_b >= 0 x_b = torch.ones((bench_tokens, bench_hidden), dtype=torch.bfloat16, device="cuda") * float(rank) - # Drive the benchmark through the high-level MoECommunicator (the public - # #818 API), mode=HIGH_THROUGHPUT. It owns its own ExpertParallelRuntime - # sized for the bench shape and runs get_dispatch_layout + intranode - # dispatch/combine internally. The first (uncached) dispatch records the - # routing layout on the returned handle; subsequent dispatches reuse it via - # previous_handle, skipping notify_dispatch's host-side counter wait. This - # isolates the on-GPU dispatch-kernel cost (NCCL-EP ep_bench convention). + # Drive the benchmark through the public high-level API. The first + # (uncached) dispatch records the routing layout on the returned handle; + # subsequent dispatches reuse it via previous_handle, skipping notify's + # host-side counter wait. This isolates the on-GPU dispatch-kernel cost + # (NCCL-EP ep_bench convention). moe = ep.MoECommunicator( - group=group, + comm=ep_group, num_experts=bench_num_experts, hidden_size=bench_hidden, topk=bench_num_topk, diff --git a/test/python/ext/ep/test_low_latency_multirank.py b/test/python/ep/test_low_latency_multirank.py similarity index 95% rename from test/python/ext/ep/test_low_latency_multirank.py rename to test/python/ep/test_low_latency_multirank.py index 72c7efcd..a00ae1b1 100644 --- a/test/python/ext/ep/test_low_latency_multirank.py +++ b/test/python/ep/test_low_latency_multirank.py @@ -3,18 +3,18 @@ """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/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: MASTER_ADDR= MASTER_PORT=29600 NODE_RANK=0 \ torchrun --nnodes=2 --nproc_per_node=1 --rdzv-backend=c10d \ - --rdzv-endpoint=:29600 test/python/ext/ep/test_low_latency_multirank.py + --rdzv-endpoint=:29600 test/python/ep/test_low_latency_multirank.py # node 1: MASTER_ADDR= MASTER_PORT=29600 NODE_RANK=1 \ torchrun --nnodes=2 --nproc_per_node=1 --rdzv-backend=c10d \ - --rdzv-endpoint=:29600 test/python/ext/ep/test_low_latency_multirank.py + --rdzv-endpoint=:29600 test/python/ep/test_low_latency_multirank.py Exercises the LL dispatch + combine round-trip on a single node. The minimal correctness check: @@ -79,7 +79,7 @@ def main(): args = parse_args() rank, num_ranks, local_rank, group = init_dist() from mscclpp import CommGroup - from mscclpp.ext import ep + import mscclpp.ep as ep ep_group = CommGroup(torch_group=group) @@ -149,8 +149,9 @@ def main(): output_buffer=dispatch_output_buffer, ) packed_recv_x = dispatch_out.tokens - packed_recv_count = dispatch_out.num_tokens_per_expert - packed_recv_layout_range = handle.layout_range + assert dispatch_out.layout.num_tokens_per_expert is not None + packed_recv_count = dispatch_out.layout.num_tokens_per_expert + packed_recv_layout_range = handle.combine_context.layout_range torch.cuda.synchronize() print(f"[rank {rank}] post-dispatch", flush=True) # packed_recv_x: [num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden] @@ -258,7 +259,8 @@ def main(): end_ev.record() torch.cuda.synchronize() disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters - recv_tokens = int(dout[0].num_tokens_per_expert.sum().item()) + assert dout[0].layout.num_tokens_per_expert is not None + recv_tokens = int(dout[0].layout.num_tokens_per_expert.sum().item()) dist.barrier(group=group) start_ev.record()