microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-13 10:47:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

Binyli/ep revise (#828)

Browse Source 
This pull request makes significant improvements to the MoE (Mixture of
Experts) Python API and documentation, focusing on clarifying and
expanding the Expert Parallel (EP) interface, especially around
quantization, dispatch/combine handles, and overlap configuration. The
changes introduce new data structures, update function signatures, and
improve documentation to better reflect the current and planned
capabilities of the system. Additionally, the base development container
is updated to CUDA 13.0, and minor corrections are made to extension
naming.
This commit is contained in:
Binyang Li
2026-07-06 21:14:29 -07:00
committed by GitHub
parent c398af350d
commit 8e34326d7a
19 changed files with 1585 additions and 1358 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/csrc/core_py.cpp
View File

@@ -62,20 +62,20 @@ void register_core(nb::module_& m) {
void* data = reinterpret_cast<void*>(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<nb::gil_scoped_release>())
.def(
"recv",
[](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) {
void* data = reinterpret_cast<void*>(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<nb::gil_scoped_release>())
.def("all_gather", &Bootstrap::allGather, nb::arg("allData"), nb::arg("size"))
.def("barrier", &Bootstrap::barrier)
.def("send", static_cast<void (Bootstrap::*)(const std::vector<char>&, 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<nb::gil_scoped_release>())
.def("recv", static_cast<void (Bootstrap::*)(std::vector<char>&, int, int)>(&Bootstrap::recv), nb::arg("data"),
nb::arg("peer"), nb::arg("tag"));
nb::arg("peer"), nb::arg("tag"), nb::call_guard<nb::gil_scoped_release>());
nb::class_<UniqueId>(m, "CppUniqueId")
.def(nb::init<>())

242
python/mscclpp/ext/ep/README.md → python/mscclpp/ep/README.md
View File

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

54
python/mscclpp/ep/__init__.py Normal file
View File

@@ -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",
]

23
python/mscclpp/ep/_cpp.py Normal file
View File

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

170
python/mscclpp/ep/communicator.py Normal file
View File

@@ -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

470
python/mscclpp/ext/ep/buffer.py → python/mscclpp/ep/high_throughput.py
View File

@@ -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), "<u2", owner), device="cuda")
return u16.view(torch.bfloat16)
class ExpertParallelRuntime:
class HighThroughputRuntime:
"""Core high-throughput expert-parallel (EP) communication runtime.
``group`` is the ``torch.distributed`` process group used only for the
``comm`` is the ``mscclpp.CommGroup`` used for rank information and
out-of-band exchange of device ids, CUDA-IPC handles, and the MSCCL++ unique
id. All dispatch/combine data movement happens through the MSCCL++ runtime.
"""
@@ -105,7 +74,7 @@ class ExpertParallelRuntime:
def __init__(
self,
group: dist.ProcessGroup,
comm: Any,
num_nvl_bytes: int = 0,
num_rdma_bytes: int = 0,
low_latency_mode: bool = False,
@@ -113,14 +82,14 @@ class ExpertParallelRuntime:
) -> 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")

332
python/mscclpp/ep/low_latency.py Normal file
View File

@@ -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}")

205
python/mscclpp/ep/types.py Normal file
View File

@@ -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"

148
python/mscclpp/ep/utils.py Normal file
View File

@@ -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), "<u2", owner), device="cuda")
return u16.view(torch.bfloat16)
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 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

6
python/mscclpp/ext/__init__.py
View File

@@ -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

38
python/mscclpp/ext/ep/__init__.py
View File

@@ -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",
]

895
python/mscclpp/ext/ep/communicator.py
View File

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