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()