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

MoE Commnucator design doc (#818)

Browse Source 
Add API doc for MoE communication
Refactor EP API for Low latency mode
This commit is contained in:
Binyang Li
2026-06-29 13:32:51 -07:00
committed by GitHub
parent 462ab1661a
commit 57ea3dd5c9
35 changed files with 3863 additions and 1718 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/mscclpp/_core/comm.py
View File

@@ -54,14 +54,21 @@ class CommGroup:
import torch
import torch.distributed as dist
backend = str(dist.get_backend(torch_group)).lower()
device = torch.device("cuda", torch.cuda.current_device()) if "nccl" in backend else torch.device("cpu")
if rank == 0:
uniq_id_global = uniq_id
pickled_data = pickle.dumps(uniq_id)
data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone()
size_tensor = torch.tensor([len(pickled_data)], dtype=torch.int64, device=device)
else:
data_tensor = torch.zeros(256, dtype=torch.uint8)
size_tensor = torch.zeros(1, dtype=torch.int64, device=device)
dist.broadcast(size_tensor, src=0, group=torch_group)
payload_size = int(size_tensor.item())
if rank == 0:
data_tensor = torch.frombuffer(bytearray(pickled_data), dtype=torch.uint8).clone().to(device)
else:
data_tensor = torch.zeros(payload_size, dtype=torch.uint8, device=device)
dist.broadcast(data_tensor, src=0, group=torch_group)
uniq_id_global = pickle.loads(data_tensor.numpy().tobytes())
uniq_id_global = pickle.loads(data_tensor.cpu().numpy().tobytes())
self.bootstrap.initialize(uniq_id_global)
elif not interfaceIpPortTrio == "":
assert rank >= 0 and size >= 1

756
python/mscclpp/ext/ep/README.md Normal file
View File

@@ -0,0 +1,756 @@
# Expert Parallel Python API Design
This document proposes a simplified Python API for MoE expert-parallel
dispatch and combine. It is a design draft for review, not a committed API
contract.
## Goals
The API should expose the tensors that the model naturally owns:
- token activations,
- top-k expert ids,
- routing weights,
- quantization scales.
It should adapt to different MoE runner backends, such as Triton ragged/grouped
GEMM, CUTLASS-style grouped GEMM, DeepGEMM, FlashInfer/CuteDSL, and custom MLP
kernels, without forcing one physical layout on every backend.
The dispatch output should make the local MLP contract explicit:
- the token layout returned by dispatch,
- per-local-expert valid token counts,
- optional expert offsets for flat layouts,
- optional quantization scale layout,
- optional overlap capability when the selected runner can notify combine.
## Public class name
Use `MoECommunicator` as the public class name:
```python
from mscclpp.ext.ep import MoECommunicator
moe_comm = MoECommunicator(...)
```
The class owns MoE dispatch/combine communication, but it does not own the MLP
compute backend.
## MoECommunicator configuration
`MoECommunicator` owns communication setup, scratch buffers, expert placement,
layout choice, and optional overlap resources. These fields should be configured
once instead of being passed to every dispatch/combine call.
```python
@dataclass
class MoECommunicatorConfig:
# Communication
comm: Optional[mscclpp.CommGroup] = None
device: Optional[torch.device | int] = None
# Expert topology
num_experts: int = 0
num_local_experts: Optional[int] = None # inferred for even contiguous placement
local_expert_start: Optional[int] = None # inferred as rank * num_local_experts
# Model shape and capacity
hidden_size: int = 0
topk: int = 0
max_tokens_per_rank: int = 0
max_recv_tokens_per_rank: Optional[int] = None
# Runtime mode and output layout
mode: MoEMode = MoEMode.LOW_LATENCY
output_layout: Optional[DispatchLayout] = None # default is derived from mode
# Quantization defaults
input_dtype: Optional[torch.dtype] = None
quant_format: Optional[str] = None
# Transport resources
num_rdma_qps_per_rank: int = 12 # RDMA QPs per peer rank; advanced tuning
num_sms: int = 20
# Overlap
enable_overlap: bool = False
```
The constructor can accept either a config object or keyword arguments:
```python
moe_comm = MoECommunicator(
comm=comm,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
topk=topk,
max_tokens_per_rank=max_tokens,
mode=MoEMode.HIGH_THROUGHPUT,
)
```
### Communication fields
`comm` is the MSCCL++ communication object used for rank information,
out-of-band metadata exchange, and connection setup.
The class should cache:
| Field | Purpose |
|---|---|
| `comm` | MSCCL++ communicator or `CommGroup` |
| `rank`, `world_size` | global EP rank information |
| `local_rank`, `device` | CUDA device binding |
| internal runtime | nanobind/C++ EP runtime implementation |
### Expert placement fields
The class needs enough information to map global expert ids to local expert
ids:
```python
local_expert_id = global_expert_id - local_expert_start
```
For the common contiguous placement:
```python
if num_local_experts is None:
assert num_experts % world_size == 0
num_local_experts = num_experts // world_size
if local_expert_start is None:
local_expert_start = rank * num_local_experts
```
So both fields may be `None` for the common even contiguous placement. If the
placement is uneven or non-contiguous, the caller must provide enough placement
information explicitly. The first version can require contiguous local experts;
a later version can add an explicit `expert_map` for arbitrary placement.
### Runtime fields
`MoECommunicator` should keep these runtime decisions internally:
| Field | Purpose |
|---|---|
| `mode` | Backend selection (`"ll"` active; `"ht"` archived/not compiled) |
| `output_layout` | MLP input layout returned by dispatch |
| `max_tokens_per_rank` | dispatch capacity |
| `max_recv_tokens_per_rank` | recv buffer capacity |
| scratch buffers | internally sized from mode, capacity, topology, and shape |
| `num_rdma_qps_per_rank`, `num_sms` | backend launch/resource tuning |
| `dispatch_config`, `combine_config` | backend-specific tuning configs |
| `overlap_capability` | whether selected MLP/backend supports notify |
The user should not pass these to `dispatch` unless explicitly overriding a
specialized advanced path.
### Mode selection
The active implementation supports `mode=MoEMode.LOW_LATENCY`. `mode` must be a
`MoEMode` enum value, not a string. `MoEMode.HIGH_THROUGHPUT` raises
`NotImplementedError` because the HT implementation is archived under
`src/ext/ep/ht/` and is not compiled into `mscclpp_ep_cpp`.
```python
moe_comm = MoECommunicator(..., mode=MoEMode.LOW_LATENCY)
```
This keeps `MoECommunicator` policy-free. Serving frameworks such as SGLang can
choose a mode based on their own scheduling policy, batch shape, runner backend,
and benchmarking data once multiple active backends are available.
The selected mode determines the default dispatch output layout:
| Mode | Default layout |
|---|---|
| `ht` | `DispatchLayout.FLAT` |
| `ll` | `DispatchLayout.EXPERT_MAJOR` |
`output_layout` may still be kept as an advanced override if a backend supports
multiple layouts within the same mode.
Use `DispatchLayout` instead of string literals for this field:
| Layout enum | Tensor shape |
|---|---|
| `DispatchLayout.FLAT` | HT: `[total_recv_tokens, hidden]`; LL: `[num_local_experts * max_slots_per_expert, hidden]` |
| `DispatchLayout.EXPERT_MAJOR` | `[num_local_experts, max_slots_per_expert, hidden]` |
## MoECommunicator methods
```python
class MoECommunicator:
def dispatch(
self,
input: torch.Tensor,
topk_ids: torch.Tensor,
weights: Optional[torch.Tensor] = None,
scales: Optional[QuantScales] = None,
*,
output_buffer: torch.Tensor,
stream: Optional[torch.cuda.Stream] = None,
) -> tuple[DispatchOutput, DispatchHandle]:
...
def combine(
self,
expert_output: torch.Tensor,
handle: DispatchHandle,
*,
out: Optional[torch.Tensor] = None,
stream: Optional[torch.cuda.Stream] = None,
) -> torch.Tensor:
...
def dispatch_async(..., overlap_config: Optional[CommOverlapConfig] = None) -> DispatchRequest:
...
def combine_async(..., overlap_config: Optional[CommOverlapConfig] = None) -> CombineRequest:
...
def create_overlap_config(
self,
op: str, # "dispatch" or "combine"
*,
handle: Optional[DispatchHandle] = None,
level: str = "op", # "op" or "block"
) -> CommOverlapConfig:
...
```
The blocking `dispatch` and `combine` methods should be the default path. The
`*_async` methods and `create_overlap_config` are optional advanced APIs for
communication/computation overlap.
If `stream` is not provided, both methods launch on `torch.cuda.current_stream()`.
## High-level API
```python
dispatch_out, handle = moe_comm.dispatch(
input,
topk_ids,
weights=None,
scales=None,
output_buffer=output_buffer,
)
expert_output = mlp(
dispatch_out.tokens,
dispatch_out.num_tokens_per_expert,
dispatch_out.scales,
)
output = moe_comm.combine(expert_output, handle)
```
`dispatch_out` is for the local MLP. `handle` is for `combine`. The MLP should
not need to inspect the opaque handle.
## Proposed types
```python
@dataclass
class QuantScales:
local: Optional[torch.Tensor] = None
global_scale: Optional[torch.Tensor] = None
format: Optional[str] = None
block_size: Optional[int] = None
class DispatchLayout(str, Enum):
FLAT = "flat"
EXPERT_MAJOR = "expert_major"
@dataclass
class DispatchOutput:
tokens: torch.Tensor
scales: Optional[QuantScales]
num_tokens_per_expert: torch.Tensor | list[int]
expert_offsets: Optional[torch.Tensor] = None
layout: DispatchLayout = DispatchLayout.FLAT
class DispatchHandle:
"""Opaque handle returned by dispatch and consumed by combine."""
@dataclass
class CommOverlapConfig:
op: str # "dispatch" or "combine"
level: str = "op" # "op" or "block"
stream: Optional[torch.cuda.Stream] = None
wait_event: Optional[torch.cuda.Event] = None
signal: Optional[torch.Tensor] = None
num_comm_sms: Optional[int] = None
block_m: Optional[int] = None
block_ready_value: Optional[int] = None
```
`create_overlap_config` creates optional overlap configuration for async
dispatch/combine calls.
```python
dispatch_overlap_config = moe_comm.create_overlap_config(op="dispatch")
combine_overlap_config = moe_comm.create_overlap_config(op="combine", handle=handle)
```
Operation-level overlap does not require `create_overlap_config`; `dispatch_async`
and `combine_async` can use their default stream/event behavior. Use
`create_overlap_config` when the caller wants explicit stream/event/SM settings
or block-level combine overlap.
For block-level MLP/combine overlap, the config includes the combine-side wait
protocol and the device signal that an overlap-capable MLP backend must publish:
```python
combine_overlap_config = moe_comm.create_overlap_config(
op="combine",
handle=handle,
level="block",
)
```
`op="dispatch", level="block"` is not part of the first version. Dispatch
overlap is operation-level only.
`CommOverlapConfig` fields:
| Field | Purpose |
|---|---|
| `op` | `"dispatch"` or `"combine"` |
| `level` | `"op"` or `"block"` |
| `stream` | Optional communication stream |
| `wait_event` | Optional event the communication op waits on before starting |
| `signal` | Device tensor written by MLP and waited on by combine for block overlap |
| `num_comm_sms` | Optional SM budget for communication |
| `block_m` | Rows per block for block overlap |
| `block_ready_value` | Signal value that marks one block as ready for combine |
`DispatchHandle` should store the metadata needed to reverse dispatch:
- source rank and source token index,
- top-k slot or equivalent routing metadata,
- top-k ids and routing weights, or stable references/copies,
- dispatch layout/range/count metadata,
- capacity, local expert placement, and launch parameters needed by kernels,
- optional cached metadata for repeated routing.
## Dispatch inputs
### `input`
`input` is the original local token matrix.
```python
input: torch.Tensor # [num_tokens, hidden]
```
Requirements:
| Field | Requirement |
|---|---|
| Shape | `[T, H]`, token-major |
| Layout | contiguous row-major |
| Device | CUDA tensor |
| dtype | BF16, FP16, FP8, NVFP4, MXFP8, or another supported activation dtype |
| Ordering | original local token order; not expert sorted |
The user should not expand `input` by top-k and should not convert it to
expert-major before calling `dispatch`.
`dispatch` includes any metadata exchange needed before moving token payloads.
For normal/high-throughput modes this typically means computing send counts from
`topk_ids`, exchanging counts or layout information across ranks, choosing recv
slots, and then dispatching the activation payload. Users should not call a
separate metadata-exchange API in the simple path.
### `topk_ids`
```python
topk_ids: torch.Tensor # [T, K], int64
```
`topk_ids[t, k]` is the global expert id selected for token `t` at top-k slot
`k`. Invalid slots may use `-1` if supported by the backend.
### `weights`
```python
weights: Optional[torch.Tensor] # [T, K], usually float32
```
These are MoE routing weights, not quantization scales. They are used by
combine to reduce the `K` expert results for each token back to `[T, H]`.
### `scales`
`scales` contains activation quantization metadata for `input`. It should be
`None` for BF16/FP16 input.
Examples:
| Format | `input` | `scales.local` | `scales.global_scale` |
|---|---|---|---|
| BF16/FP16 | `[T, H]` | `None` | `None` |
| FP8 E4M3 | `[T, H]` FP8 | `[T, H / block_size]`, often block size 128 | usually `None` |
| NVFP4 | backend-defined packed/logical `[T, H]` | block scale tensor | optional global scale |
| MXFP8 | backend-defined `[T, H]` | micro-scale tensor, e.g. E8M0 blocks | optional/global if required |
The API should not assume quantization scale is a scalar. For FP8 paths in
DeepEP/SGLang, scales are usually per token and per hidden block.
### `output_buffer`
Low-latency dispatch requires the caller to provide the receive token buffer:
```python
output_buffer: torch.Tensor
```
For padded expert-major LL layout:
```text
output_buffer: [num_local_experts, world_size * max_tokens_per_rank, hidden]
```
The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16.
For FP8 dispatch it is FP8 and the returned `DispatchOutput.scales` carries the
matching scale tensor.
`output_buffer` is required for LL because the MLP runner often owns or reuses
workspace memory. `MoECommunicator` writes dispatch output into the provided
buffer instead of allocating it internally.
## Dispatch output layout for MLP
`dispatch` should return MLP-ready tokens. The MLP should not run another
token-major to expert-major permutation unless it uses a custom adapter.
### Normal / high-throughput flat layout
HT uses `DispatchLayout.FLAT`, a flat expert-major layout:
```python
dispatch_out.tokens # [total_recv_tokens, H]
```
Rows are grouped by local expert id:
```text
expert0 tokens
expert1 tokens
expert2 tokens
...
```
`dispatch_out.num_tokens_per_expert` is ordered by local expert id:
```python
num_tokens_per_expert[i] = valid token count for local expert i
```
For flat layout, `expert_offsets` may be provided or derived by cumulative sum:
```python
expert_offsets = cumsum([0] + num_tokens_per_expert)
tokens[expert_offsets[i] : expert_offsets[i + 1]]
```
This layout is efficient for Triton or grouped GEMM kernels because it avoids
padding.
### Low-latency output layouts
LL defaults to `DispatchLayout.EXPERT_MAJOR`, a padded expert-major tensor:
```python
dispatch_out.tokens # [num_local_experts, max_slots_per_expert, H]
```
LL can also return `DispatchLayout.FLAT`, which is the same contiguous
local-expert-major storage viewed as 2D:
```python
dispatch_out.tokens # [num_local_experts * max_slots_per_expert, H]
```
For expert `i`, only the first `num_tokens_per_expert[i]` slots are valid:
```python
expert_major_tokens = dispatch_out.tokens.view(num_local_experts, max_slots_per_expert, H)
expert_major_tokens[i, :num_tokens_per_expert[i], :]
```
The remaining slots are padding or scratch space. The MLP output must keep the
same layout and slot order.
### Scale output layout
If `dispatch_out.scales` is not `None`, its local scale tensor should follow
the same packed/expert-major layout as `dispatch_out.tokens`, with the hidden
dimension replaced by the scale dimension.
Examples:
```text
flat tokens: HT [total_recv_tokens, H]; LL [num_local_experts * max_slots, H]
flat FP8 scales: HT [total_recv_tokens, H / 128]; LL [num_local_experts, max_slots, H / 128]
expert-major tokens: [num_local_experts, max_slots, H]
expert-major scales: [num_local_experts, max_slots, H / 128]
```
## MLP contract
The MLP consumes `dispatch_out`, not the original token-major input.
For flat expert-major output:
```python
expert_output = triton_mlp(
dispatch_out.tokens,
dispatch_out.num_tokens_per_expert,
dispatch_out.scales,
)
```
For padded expert-major output:
```python
expert_output = expert_major_mlp(
dispatch_out.tokens,
dispatch_out.num_tokens_per_expert,
dispatch_out.scales,
)
```
The MLP must preserve the dispatch output layout and row/slot order. It may
apply expert-specific GEMMs, but it must not compact or reorder tokens unless it
also produces compatible metadata for combine.
## Combine API
```python
output = moe_comm.combine(
expert_output,
handle,
out=None,
)
```
`expert_output` must have the same physical layout and order as
`dispatch_out.tokens`.
`combine` uses `handle` to:
- map each expert output row/slot back to the original source rank and token,
- find the corresponding top-k slot,
- apply the routing weight,
- reduce all selected expert outputs for each token,
- return local output in original token-major order.
The output shape is:
```python
output # [T, H]
```
`combine` should not require users to pass `topk_ids`, `weights`, prefix
matrices, source indices, or layout ranges again. Those belong in `handle`.
An optional `weights` override could be added later, but the default API should
use the weights captured by `dispatch`.
## Communication/computation overlap
The default API should be blocking and simple:
```python
dispatch_out, handle = moe_comm.dispatch(
input,
topk_ids,
weights,
scales,
output_buffer=output_buffer,
)
expert_output = mlp(dispatch_out.tokens, dispatch_out.num_tokens_per_expert)
output = moe_comm.combine(expert_output, handle)
```
For overlap, expose two optional APIs rather than adding many flags to the
default path:
| API | Purpose |
|---|---|
| `dispatch_async` / `combine_async` | Coarse-grained async launch and wait |
| `create_overlap_config(..., level="block")` | Fine-grained block notify between MLP down-GEMM and combine |
### Coarse-grained overlap
Coarse-grained overlap lets the caller launch communication on a communication
stream and wait later.
```python
dispatch_overlap_config = moe_comm.create_overlap_config(op="dispatch")
dispatch_req = moe_comm.dispatch_async(
input,
topk_ids,
weights,
scales,
output_buffer=output_buffer,
overlap_config=dispatch_overlap_config,
)
# Run unrelated work while dispatch metadata/payload communication is in flight.
dispatch_out, handle = dispatch_req.wait()
expert_output = mlp(dispatch_out.tokens, dispatch_out.num_tokens_per_expert)
combine_overlap_config = moe_comm.create_overlap_config(op="combine", handle=handle)
combine_req = moe_comm.combine_async(
expert_output,
handle,
overlap_config=combine_overlap_config,
)
# Run unrelated work while combine is in flight.
output = combine_req.wait()
```
This is similar to the event/hook style used by DeepEP and SGLang. The request
object should own any stream event or receive hook required by the backend.
### Fine-grained MLP/combine overlap
Fine-grained overlap sends combine data as soon as the MLP produces a block.
This requires a device-side notify/signal from the MLP backend to the combine
kernel.
```python
combine_overlap_config = moe_comm.create_overlap_config(
op="combine",
handle=handle,
level="block",
)
# User must adapt the MLP backend/adapter to consume this config and notify
# combine as blocks become ready.
config = combine_overlap_config
expert_output = mlp(
dispatch_out.tokens,
dispatch_out.num_tokens_per_expert,
config=config,
)
combine_req = moe_comm.combine_async(
expert_output,
handle,
overlap_config=combine_overlap_config,
)
output = combine_req.wait()
```
The overlap config is not routing metadata. It only tells combine when a
region of `expert_output` is ready to read. The routing/source mapping still
comes from `handle`.
The MLP backend must follow these rules when using notify:
- write `expert_output` in the same row/slot order as `dispatch_out.tokens`,
- publish data before signaling readiness,
- signal at the block granularity defined by `overlap_config`,
- use the signal value/protocol provided by `overlap_config`.
If the MLP backend does not support notify, it can still use the blocking API or
coarse-grained `combine_async` after the full `expert_output` tensor is ready.
This must be a joint contract between the dispatcher and the MLP runner. The
dispatcher can provide the signal buffer and combine-side wait protocol, but it
cannot infer readiness by itself. The MLP runner must write the signal after it
finishes the corresponding output region.
SGLang follows this model for its DeepEP low-latency path. It computes overlap
arguments after dispatch, passes combine-side arguments to the DeepEP dispatcher,
and passes down-GEMM arguments to the MoE runner. Backend support is selective:
- DeepGEMM FP8 masked down-GEMM can return block metadata such as `block_m` and
`block_ready_value` and signal combine readiness.
- FlashInfer CuteDSL can receive down-GEMM signal/start-event arguments.
- Some paths, such as BF16 masked DeepGEMM and generic Triton runners, do not
support this block overlap protocol.
Therefore, the API should expose overlap as an optional capability advertised by
the MLP backend, not as a guaranteed feature of every `combine_async` call.
## Internal metadata exchange
Normal/high-throughput dispatch usually needs a metadata phase before payload
movement:
```text
topk_ids
-> compute send counts per rank/expert
-> exchange counts or layout metadata
-> compute recv slots and local expert counts
-> dispatch token payload
```
Low-latency modes may use fixed-capacity buffers and device-side counters, but
they still generate metadata such as source info, layout ranges, and valid
counts.
These details should remain internal. The user-facing API should only expose
MLP-relevant layout information through `DispatchOutput` and combine-relevant
metadata through `DispatchHandle`.
## Example
```python
recv, handle = moe_comm.dispatch(
input=hidden_states, # [T, H]
topk_ids=topk_ids, # [T, K]
weights=topk_weights, # [T, K]
scales=None, # BF16 path
output_buffer=recv_buffer,
)
expert_output = triton_grouped_mlp(
recv.tokens,
recv.num_tokens_per_expert,
)
output = moe_comm.combine(expert_output, handle)
```
Quantized path:
```python
recv, handle = moe_comm.dispatch(
input=x_fp8,
topk_ids=topk_ids,
weights=topk_weights,
scales=QuantScales(
local=x_scales,
format="fp8_e4m3",
block_size=128,
),
output_buffer=recv_buffer,
)
expert_output = fp8_grouped_mlp(
recv.tokens,
recv.scales,
recv.num_tokens_per_expert,
)
output = moe_comm.combine(expert_output, handle)
```

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

@@ -2,12 +2,28 @@
# Licensed under the MIT License.
"""MSCCL++ Expert-Parallel (MoE dispatch/combine) extension.
See ``src/ext/ep/README.md`` in the repository for migration status. The
``Buffer`` class mirrors :class:`deep_ep.Buffer` and supports intranode
(NVLink-only) dispatch/combine as well as internode HT and low-latency
paths.
See ``src/ext/ep/README.md`` in the repository for migration status.
``MoECommunicator`` is the high-level public API.
"""
from .buffer import Buffer, Config, EventHandle # noqa: F401
from .communicator import ( # noqa: F401
CommOverlapConfig,
DispatchHandle,
DispatchOutput,
DispatchLayout,
MoEMode,
MoECommunicator,
MoECommunicatorConfig,
QuantScales,
)
__all__ = ["Buffer", "Config", "EventHandle"]
__all__ = [
"CommOverlapConfig",
"DispatchHandle",
"DispatchOutput",
"DispatchLayout",
"MoEMode",
"MoECommunicator",
"MoECommunicatorConfig",
"QuantScales",
]

193
python/mscclpp/ext/ep/buffer.py
View File

@@ -1,193 +0,0 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
#
# Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP),
# branch ``chhwang/dev-atomic-add-cleanup``. Licensed under the MIT License.
"""Python frontend for the MSCCL++ Expert-Parallel extension.
This is a thin wrapper around the pybind11 extension ``mscclpp_ep_cpp``.
The shape of :class:`Buffer` mirrors :class:`deep_ep.Buffer` so existing
DeepEP users can port with minimal changes.
Current status (see ``src/ext/ep/README.md``):
* Intranode (NVLink-only) dispatch and combine: ported and validated on
one node with 8 GPUs.
* ``get_dispatch_layout``: ported.
* Internode HT (MSCCL++ PortChannel + MemoryChannel) dispatch and combine:
ported and validated on 2 nodes x 8 H100 GPUs with
``test/python/ext/ep/test_internode_multirank.py``.
* Internode low-latency kernels (NVSHMEM/IBGDA -> MSCCL++ PortChannel):
ported and validated on 2 nodes x 8 H100 GPUs with
``test/python/ext/ep/test_low_latency_multirank.py``.
"""
from __future__ import annotations
import os
from typing import List, Optional, Tuple
import torch
import torch.distributed as dist
try:
import mscclpp_ep_cpp as _cpp # type: ignore[import-not-found]
except ImportError as exc: # pragma: no cover
raise ImportError(
"mscclpp_ep_cpp is not available. Build mscclpp with "
"-DMSCCLPP_BUILD_EXT_EP=ON (and ensure PyTorch's CMake prefix is on "
"CMAKE_PREFIX_PATH) or install via `pip install` after the build."
) from exc
Config = _cpp.Config
EventHandle = _cpp.EventHandle
class Buffer:
"""Core expert-parallel (EP) communication buffer.
Parameters
----------
group:
The ``torch.distributed`` process group. Used only for out-of-band
exchange of IPC handles and the MSCCL++ unique id.
num_nvl_bytes:
Size of the NVLink-accessible scratch buffer (shared via CUDA IPC).
num_rdma_bytes:
Size of the RDMA scratch buffer. Required (>0) for internode HT and
low-latency modes.
low_latency_mode:
Enable the low-latency dispatch/combine path. This mode uses only
the RDMA buffer (``num_rdma_bytes``) and drives every peer through
MSCCL++ ``PortChannel``; consequently, it works cross-node with any
topology but is still pending H100 hardware validation.
num_qps_per_rank:
Ignored for intranode mode.
"""
#: Default number of SMs reserved for comms kernels. Matches DeepEP.
num_sms: int = 20
def __init__(
self,
group: dist.ProcessGroup,
num_nvl_bytes: int = 0,
num_rdma_bytes: int = 0,
low_latency_mode: bool = False,
num_qps_per_rank: int = 12,
) -> None:
self.rank: int = group.rank()
self.group_size: int = group.size()
self.group = group
self.num_nvl_bytes = num_nvl_bytes
self.num_rdma_bytes = num_rdma_bytes
self.low_latency_mode = low_latency_mode
self.num_qps_per_rank = num_qps_per_rank
self.runtime = _cpp.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode)
# Exchange device IDs + IPC handles + (for RDMA) the MSCCL++ unique id.
device_ids: List[Optional[int]] = [None] * self.group_size
local_device_id = self.runtime.get_local_device_id()
dist.all_gather_object(device_ids, local_device_id, group)
ipc_handles: List[Optional[bytes]] = [None] * self.group_size
local_ipc_handle = self.runtime.get_local_ipc_handle()
dist.all_gather_object(ipc_handles, local_ipc_handle, group)
root_unique_id: Optional[bytes] = None
# MSCCL++ requires a bootstrapped Communicator even for pure-NVLink
# setups because `Buffer::sync()` uses `communicator->connect(ipc)`
# to build MemoryChannels. We always exchange a unique id.
if num_qps_per_rank <= 0:
raise ValueError("num_qps_per_rank must be > 0")
if self.rank == 0:
root_unique_id = self.runtime.create_unique_id()
broadcast_list = [root_unique_id]
dist.broadcast_object_list(broadcast_list, src=0, group=group)
root_unique_id = broadcast_list[0]
assert root_unique_id is not None
self.runtime.connect(root_unique_id)
# sync() expects Sequence[bytearray | None] / bytearray | None.
ipc_handles_ba = [bytearray(h) if h is not None else None for h in ipc_handles]
self.runtime.sync(device_ids, ipc_handles_ba, bytearray(root_unique_id))
# ------------------------------------------------------------------
# Sanity helpers
# ------------------------------------------------------------------
def is_available(self) -> bool:
return self.runtime.is_available()
def is_internode_available(self) -> bool:
return self.runtime.is_internode_available()
def get_local_device_id(self) -> int:
return self.runtime.get_local_device_id()
def get_num_rdma_ranks(self) -> int:
return self.runtime.get_num_rdma_ranks()
def get_rdma_rank(self) -> int:
return self.runtime.get_rdma_rank()
def get_root_rdma_rank(self, global_: bool) -> int:
return self.runtime.get_root_rdma_rank(global_)
# ------------------------------------------------------------------
# Layout / dispatch / combine (thin pass-through wrappers).
# Signatures mirror deep_ep.Buffer so existing test harnesses can reuse.
# ------------------------------------------------------------------
def get_dispatch_layout(
self,
topk_idx: torch.Tensor,
num_experts: int,
previous_event: Optional[EventHandle] = None,
async_finish: bool = False,
allocate_on_comm_stream: bool = False,
):
return self.runtime.get_dispatch_layout(
topk_idx, num_experts, previous_event, async_finish, allocate_on_comm_stream
)
def intranode_dispatch(self, *args, **kwargs):
return self.runtime.intranode_dispatch(*args, **kwargs)
def intranode_combine(self, *args, **kwargs):
return self.runtime.intranode_combine(*args, **kwargs)
def internode_dispatch(self, *args, **kwargs):
return self.runtime.internode_dispatch(*args, **kwargs)
def internode_combine(self, *args, **kwargs):
return self.runtime.internode_combine(*args, **kwargs)
def clean_low_latency_buffer(self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int) -> None:
self.runtime.clean_low_latency_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
def low_latency_dispatch(self, *args, **kwargs):
return self.runtime.low_latency_dispatch(*args, **kwargs)
def low_latency_combine(self, *args, **kwargs):
return self.runtime.low_latency_combine(*args, **kwargs)
def get_next_low_latency_combine_buffer(self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int):
return self.runtime.get_next_low_latency_combine_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)
def get_local_buffer_tensor(
self, dtype: torch.dtype, offset: int = 0, use_rdma_buffer: bool = False
) -> torch.Tensor:
return self.runtime.get_local_buffer_tensor(dtype, offset, use_rdma_buffer)
# ------------------------------------------------------------------
# Static helpers
# ------------------------------------------------------------------
@staticmethod
def get_low_latency_rdma_size_hint(
num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int
) -> int:
return _cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)

536
python/mscclpp/ext/ep/communicator.py Normal file
View File

@@ -0,0 +1,536 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
#
# Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP),
# branch ``chhwang/dev-atomic-add-cleanup``. Licensed under the MIT License.
"""Python frontend for the MSCCL++ Expert-Parallel extension.
This is a thin wrapper around the nanobind extension ``mscclpp_ep_cpp``.
``MoECommunicator`` is the high-level API. ``_MoERuntime`` is a private
low-latency runtime wrapper used internally by the high-level API.
Current status (see ``src/ext/ep/README.md``):
* Intranode (NVLink-only) dispatch and combine: ported and validated on
one node with 8 GPUs.
* ``get_dispatch_layout``: ported.
* Internode HT (MSCCL++ PortChannel + MemoryChannel) dispatch and combine:
ported and validated on 2 nodes x 8 H100 GPUs with
``test/python/ext/ep/test_internode_multirank.py``.
* Low-latency kernels (RDMA + CUDA IPC paths):
ported and validated on intra-node and 2 nodes x 8 H100 GPUs with
``test/python/ext/ep/test_low_latency_multirank.py``.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
import torch
from mscclpp._core import CommGroup
try:
import mscclpp_ep_cpp as _cpp # type: ignore[import-not-found]
except ImportError as exc: # pragma: no cover
raise ImportError(
"mscclpp_ep_cpp is not available. Build mscclpp with -DMSCCLPP_BUILD_EXT_EP=ON "
"or install with `pip install .[ep]`."
) from exc
DispatchLayout = _cpp.DispatchLayout
MoEMode = _cpp.MoEMode
@dataclass
class MoECommunicatorConfig:
"""Configuration for the high-level MoE dispatch/combine API."""
comm: Optional[CommGroup] = None
device: Optional[torch.device | int] = None
num_experts: int = 0
num_local_experts: Optional[int] = None
local_expert_start: Optional[int] = None
hidden_size: int = 0
topk: int = 0
max_tokens_per_rank: int = 0
max_recv_tokens_per_rank: Optional[int] = None
mode: MoEMode = MoEMode.LOW_LATENCY
output_layout: Optional[DispatchLayout] = None
input_dtype: Optional[torch.dtype] = None
quant_format: Optional[str] = None
num_rdma_qps_per_rank: int = 12
num_sms: int = 20
enable_overlap: bool = False
@dataclass
class QuantScales:
local: Optional[torch.Tensor] = None
global_scale: Optional[torch.Tensor] = None
format: Optional[str] = None
block_size: Optional[int] = None
@dataclass
class DispatchOutput:
tokens: torch.Tensor
scales: Optional[QuantScales]
num_tokens_per_expert: torch.Tensor | list[int]
expert_offsets: Optional[torch.Tensor] = None
layout: DispatchLayout = DispatchLayout.FLAT
@dataclass
class DispatchHandle:
"""Opaque dispatch metadata consumed by :meth:`MoECommunicator.combine`."""
topk_ids: torch.Tensor
weights: torch.Tensor
src_info: torch.Tensor
layout_range: torch.Tensor
num_max_dispatch_tokens_per_rank: int
num_experts: int
num_tokens: int
hidden_size: int
num_local_experts: int
local_expert_start: int
layout: DispatchLayout
output_scales: Optional[QuantScales] = None
@dataclass
class CommOverlapConfig:
op: str
level: str = "op"
stream: Optional[torch.cuda.Stream] = None
wait_event: Optional[torch.cuda.Event] = None
signal: Optional[torch.Tensor] = None
num_comm_sms: Optional[int] = None
block_m: Optional[int] = None
block_ready_value: Optional[int] = None
class _MoERuntime:
"""Private low-level MoE communication runtime wrapper.
Parameters
----------
comm:
The :class:`mscclpp.CommGroup`. Used only for out-of-band
exchange of IPC handles and the MSCCL++ unique id.
num_nvl_bytes:
Size of the NVLink-accessible scratch buffer. Reserved for archived HT mode.
num_rdma_bytes:
Size of the LL RDMA scratch buffer.
mode:
Runtime mode selector. ``MoEMode.LOW_LATENCY`` is active;
``MoEMode.HIGH_THROUGHPUT`` is archived and not compiled.
num_qps_per_rank:
RDMA QPs per peer rank.
"""
#: Default number of SMs reserved for comms kernels. Matches DeepEP.
num_sms: int = 20
def __init__(
self,
comm: CommGroup,
num_nvl_bytes: int = 0,
num_rdma_bytes: int = 0,
mode: MoEMode = MoEMode.LOW_LATENCY,
num_qps_per_rank: int = 12,
) -> None:
if not isinstance(mode, MoEMode):
raise TypeError("mode must be a MoEMode")
self.mode = mode
if self.mode != MoEMode.LOW_LATENCY:
raise NotImplementedError("mode='ht' is archived under src/ext/ep/ht and is not compiled")
self.rank: int = comm.my_rank
self.group_size: int = comm.nranks
self.comm = comm
self.num_nvl_bytes = num_nvl_bytes
self.num_rdma_bytes = num_rdma_bytes
self.num_qps_per_rank = num_qps_per_rank
self._cpp_runtime = _cpp.MoERuntime(comm.communicator, num_nvl_bytes, num_rdma_bytes, mode)
if num_qps_per_rank <= 0:
raise ValueError("num_qps_per_rank must be > 0")
# ------------------------------------------------------------------
# Sanity helpers
# ------------------------------------------------------------------
def is_available(self) -> bool:
return self._cpp_runtime.is_available()
def is_internode_available(self) -> bool:
return self._cpp_runtime.is_internode_available()
def get_local_device_id(self) -> int:
return self._cpp_runtime.get_local_device_id()
def get_num_rdma_ranks(self) -> int:
return self._cpp_runtime.get_num_rdma_ranks()
def get_rdma_rank(self) -> int:
return self._cpp_runtime.get_rdma_rank()
def get_root_rdma_rank(self, global_: bool) -> int:
return self._cpp_runtime.get_root_rdma_rank(global_)
class MoECommunicator:
"""High-level MoE communicator API for dispatch/combine.
The first implementation supports the low-latency backend.
"""
def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> None:
if config is not None and kwargs:
raise ValueError("Pass either MoECommunicatorConfig or keyword arguments, not both")
if config is None:
if "group" in kwargs and "comm" not in kwargs:
kwargs["comm"] = kwargs.pop("group")
config = MoECommunicatorConfig(**kwargs)
if config.device is not None:
torch.cuda.set_device(config.device)
comm = config.comm
if comm is None:
raise ValueError("MoECommunicator requires an mscclpp.CommGroup")
self.comm = comm
self.rank: int = comm.my_rank
self.world_size: int = comm.nranks
self.local_rank: int = torch.cuda.current_device()
self.device = torch.device("cuda", self.local_rank)
if not isinstance(config.mode, MoEMode):
raise TypeError("MoECommunicatorConfig.mode must be a MoEMode")
self.mode = config.mode
if self.mode != MoEMode.LOW_LATENCY:
raise NotImplementedError("mode='ht' is archived under src/ext/ep/ht and is not compiled")
self.output_layout = _resolve_output_layout(config.output_layout, self.mode)
self.num_experts = config.num_experts
self.hidden_size = config.hidden_size
self.topk = config.topk
self.max_tokens_per_rank = config.max_tokens_per_rank
if self.num_experts <= 0 or self.hidden_size <= 0 or self.topk <= 0 or self.max_tokens_per_rank <= 0:
raise ValueError("num_experts, hidden_size, topk, and max_tokens_per_rank must be positive")
if self.num_experts % self.world_size != 0:
raise ValueError("low-latency mode requires num_experts divisible by world_size")
self.num_local_experts = config.num_local_experts
if self.num_local_experts is None:
self.num_local_experts = self.num_experts // self.world_size
if self.num_local_experts * self.world_size != self.num_experts:
raise NotImplementedError("only even contiguous expert placement is currently supported")
self.local_expert_start = config.local_expert_start
if self.local_expert_start is None:
self.local_expert_start = self.rank * self.num_local_experts
if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank):
raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity")
if config.input_dtype not in (None, torch.bfloat16):
raise NotImplementedError("low-latency dispatch currently supports BF16 input only")
self.quant_format = _normalize_quant_format(config.quant_format)
if self.quant_format not in (None, "fp8_e4m3"):
raise NotImplementedError(f"unsupported low-latency quant_format: {config.quant_format}")
self.dispatch_requires_quantization = self.quant_format is not None
num_nvl_bytes = 0
num_rdma_bytes = _get_low_latency_rdma_size_hint(
self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts
)
self.enable_overlap = config.enable_overlap
self.num_sms = config.num_sms
self._dispatch_scales: Optional[torch.Tensor] = None
self._dispatch_src_info: Optional[torch.Tensor] = None
self._dispatch_layout_range: Optional[torch.Tensor] = None
self._dispatch_count: Optional[torch.Tensor] = None
self._runtime = _MoERuntime(
comm,
num_nvl_bytes=num_nvl_bytes,
num_rdma_bytes=num_rdma_bytes,
mode=self.mode,
num_qps_per_rank=config.num_rdma_qps_per_rank,
)
def is_available(self) -> bool:
return self._runtime.is_available()
def is_internode_available(self) -> bool:
return self._runtime.is_internode_available()
def dispatch(
self,
input: torch.Tensor,
topk_ids: torch.Tensor,
weights: Optional[torch.Tensor] = None,
scales: Optional[QuantScales] = None,
*,
output_buffer: torch.Tensor,
stream: Optional[torch.cuda.Stream] = None,
) -> tuple[DispatchOutput, DispatchHandle]:
self._validate_dispatch_inputs(input, topk_ids, weights, scales, output_buffer)
if weights is None:
weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device)
output_tensors = self._get_dispatch_output_tensors(output_buffer)
output_buffer, packed_scales, src_info, layout_range, num_tokens_per_expert = output_tensors
self._runtime._cpp_runtime.dispatch(
input.data_ptr(),
topk_ids.data_ptr(),
output_buffer.data_ptr(),
0 if packed_scales is None else packed_scales.data_ptr(),
src_info.data_ptr(),
layout_range.data_ptr(),
num_tokens_per_expert.data_ptr(),
input.size(0),
self.hidden_size,
self.topk,
self.max_tokens_per_rank,
self.num_experts,
self.dispatch_requires_quantization,
self.output_layout,
_cuda_stream_ptr(stream),
)
output_scales = None
if packed_scales is not None:
output_scales = QuantScales(local=packed_scales, format="fp8_e4m3", block_size=128)
dispatch_out = DispatchOutput(
tokens=output_buffer,
scales=output_scales,
num_tokens_per_expert=num_tokens_per_expert,
expert_offsets=None,
layout=self.output_layout,
)
handle = DispatchHandle(
topk_ids=topk_ids,
weights=weights,
src_info=src_info,
layout_range=layout_range,
num_max_dispatch_tokens_per_rank=self.max_tokens_per_rank,
num_experts=self.num_experts,
num_tokens=input.size(0),
hidden_size=self.hidden_size,
num_local_experts=self.num_local_experts,
local_expert_start=self.local_expert_start,
layout=self.output_layout,
output_scales=output_scales,
)
return dispatch_out, handle
def _get_dispatch_output_tensors(self, output_buffer: torch.Tensor) -> tuple[
torch.Tensor,
Optional[torch.Tensor],
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
device = output_buffer.device
slots_per_expert = self.world_size * self.max_tokens_per_rank
if self._dispatch_src_info is None or self._dispatch_src_info.device != device:
self._dispatch_src_info = torch.empty(
(self.num_local_experts, slots_per_expert),
dtype=torch.int32,
device=device,
)
self._dispatch_layout_range = torch.empty(
(self.num_local_experts, self.world_size),
dtype=torch.int64,
device=device,
)
self._dispatch_count = torch.empty(
(self.num_local_experts,),
dtype=torch.int32,
device=device,
)
self._dispatch_scales = None
if self.dispatch_requires_quantization:
num_scales = self.hidden_size // 128
scales_storage = torch.empty(
(self.num_local_experts, num_scales, slots_per_expert),
dtype=torch.float32,
device=device,
)
self._dispatch_scales = scales_storage.transpose(1, 2)
assert self._dispatch_src_info is not None
assert self._dispatch_layout_range is not None
assert self._dispatch_count is not None
return (
output_buffer,
self._dispatch_scales,
self._dispatch_src_info,
self._dispatch_layout_range,
self._dispatch_count,
)
def combine(
self,
expert_output: torch.Tensor,
handle: DispatchHandle,
*,
out: Optional[torch.Tensor] = None,
stream: Optional[torch.cuda.Stream] = None,
) -> torch.Tensor:
self._validate_combine_inputs(expert_output, handle, out)
combine_requires_dequantization = _requires_dequantization(expert_output)
x_scales = None
if combine_requires_dequantization:
if handle.output_scales is None or handle.output_scales.local is None:
raise ValueError("FP8 expert_output requires scales captured in the dispatch handle")
x_scales = handle.output_scales.local
if out is None:
out = torch.empty((handle.num_tokens, self.hidden_size), dtype=torch.bfloat16, device=expert_output.device)
self._runtime._cpp_runtime.combine(
expert_output.data_ptr(),
0 if x_scales is None else x_scales.data_ptr(),
handle.topk_ids.data_ptr(),
handle.weights.data_ptr(),
handle.src_info.data_ptr(),
handle.layout_range.data_ptr(),
out.data_ptr(),
handle.num_tokens,
self.hidden_size,
handle.weights.size(1),
handle.num_max_dispatch_tokens_per_rank,
handle.num_experts,
combine_requires_dequantization,
_cuda_stream_ptr(stream),
)
return out
def dispatch_async(self, *args, **kwargs):
raise NotImplementedError("dispatch_async is not implemented for MoECommunicator yet")
def combine_async(self, *args, **kwargs):
raise NotImplementedError("combine_async is not implemented for MoECommunicator yet")
def create_overlap_config(
self,
op: str,
*,
handle: Optional[DispatchHandle] = None,
level: str = "op",
) -> CommOverlapConfig:
if op not in ("dispatch", "combine"):
raise ValueError("op must be 'dispatch' or 'combine'")
if level != "op":
raise NotImplementedError("block-level overlap is not implemented yet")
if op == "combine" and handle is None:
raise ValueError("combine overlap config requires a DispatchHandle")
return CommOverlapConfig(op=op, level=level)
def _validate_dispatch_inputs(
self,
input: torch.Tensor,
topk_ids: torch.Tensor,
weights: Optional[torch.Tensor],
scales: Optional[QuantScales],
output_buffer: torch.Tensor,
) -> None:
if output_buffer is None:
raise ValueError("output_buffer is required for low-latency dispatch")
if scales is not None and (scales.local is not None or scales.global_scale is not None):
raise NotImplementedError("low-latency dispatch does not support quantized input scales yet")
if input.dim() != 2 or not input.is_contiguous():
raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor")
if input.device.type != "cuda" or input.dtype != torch.bfloat16:
raise ValueError("low-latency dispatch input must be a CUDA BF16 tensor")
if input.size(1) != self.hidden_size:
raise ValueError(f"input hidden size {input.size(1)} does not match configured {self.hidden_size}")
if input.size(0) > self.max_tokens_per_rank:
raise ValueError("input token count exceeds max_tokens_per_rank")
if topk_ids.dim() != 2 or not topk_ids.is_contiguous():
raise ValueError("topk_ids must be a contiguous [num_tokens, topk] tensor")
if topk_ids.device != input.device or topk_ids.dtype != torch.int64:
raise ValueError("topk_ids must be an int64 CUDA tensor on the same device as input")
if topk_ids.shape != (input.size(0), self.topk):
raise ValueError("topk_ids shape must match [input.size(0), configured topk]")
if weights is not None:
if weights.dim() != 2 or not weights.is_contiguous():
raise ValueError("weights must be a contiguous [num_tokens, topk] tensor")
if weights.device != input.device or weights.dtype != torch.float32:
raise ValueError("weights must be a float32 CUDA tensor on the same device as input")
if weights.shape != topk_ids.shape:
raise ValueError("weights shape must match topk_ids")
expected_dtype = torch.float8_e4m3fn if self.dispatch_requires_quantization else torch.bfloat16
slots_per_expert = self.world_size * self.max_tokens_per_rank
if self.output_layout == DispatchLayout.EXPERT_MAJOR:
expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)
else:
expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size)
if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous():
raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor")
if output_buffer.device != input.device or output_buffer.dtype != expected_dtype:
raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input")
if tuple(output_buffer.shape) != expected_shape:
raise ValueError(f"output_buffer shape must be {expected_shape}")
def _validate_combine_inputs(
self, expert_output: torch.Tensor, handle: DispatchHandle, out: Optional[torch.Tensor]
) -> None:
if handle.num_experts != self.num_experts or handle.hidden_size != self.hidden_size:
raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration")
slots_per_expert = self.world_size * self.max_tokens_per_rank
if handle.layout == DispatchLayout.EXPERT_MAJOR:
expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size)
else:
expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size)
if expert_output.dim() != len(expected_shape) or not expert_output.is_contiguous():
raise ValueError("expert_output must keep dispatch output's contiguous layout")
if tuple(expert_output.shape) != expected_shape:
raise ValueError(f"expert_output shape must be {expected_shape}")
if expert_output.dtype not in (torch.bfloat16, getattr(torch, "float8_e4m3fn", None)):
raise ValueError("expert_output must be BF16 or FP8 E4M3")
if out is not None:
expected_out_shape = (handle.num_tokens, self.hidden_size)
if tuple(out.shape) != expected_out_shape or out.dtype != torch.bfloat16 or not out.is_contiguous():
raise ValueError(f"out must be a contiguous BF16 tensor with shape {expected_out_shape}")
def _resolve_output_layout(layout: Optional[DispatchLayout], mode: MoEMode) -> DispatchLayout:
if layout is None:
return DispatchLayout.EXPERT_MAJOR if mode == MoEMode.LOW_LATENCY else DispatchLayout.FLAT
if not isinstance(layout, DispatchLayout):
raise TypeError("MoECommunicatorConfig.output_layout must be a DispatchLayout")
return layout
def _cuda_stream_ptr(stream: Optional[torch.cuda.Stream]) -> int:
return (stream if stream is not None else torch.cuda.current_stream()).cuda_stream
def _normalize_quant_format(fmt: Optional[str]) -> Optional[str]:
if fmt is None:
return None
normalized = fmt.lower().replace("-", "_")
if normalized in ("fp8", "fp8_e4m3", "f8e4m3", "float8_e4m3fn"):
return "fp8_e4m3"
return normalized
def _requires_dequantization(tensor: torch.Tensor) -> bool:
fp8_dtype = getattr(torch, "float8_e4m3fn", None)
return fp8_dtype is not None and tensor.dtype == fp8_dtype
def _get_low_latency_rdma_size_hint(
num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int
) -> int:
return _cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts)

135
python/test/test_gpu_buffer_pool_nvls_zero.py Normal file
View File

@@ -0,0 +1,135 @@
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import socket
import cupy as cp
import pytest
from mscclpp import CommGroup, DataType, RawGpuBuffer, ReduceOp, GpuBufferPool, is_nvls_supported
from mscclpp.ext import AlgorithmCollectionBuilder
from mscclpp_benchmark.gpu import capture_graph
from .mscclpp_mpi import MpiGroup, parametrize_mpi_groups, mpi_group # noqa: F401
def _same_host(comm) -> bool:
hostnames = comm.allgather(socket.gethostname())
return len(set(hostnames)) == 1
def _build_nvls_zero_algorithm(mpi_group: MpiGroup):
comm_group = CommGroup(mpi_group.comm)
scratch = RawGpuBuffer(1 << 27)
AlgorithmCollectionBuilder.reset()
builder = AlgorithmCollectionBuilder()
algorithms = builder.build_default_algorithms(
scratch_buffer=scratch.data(),
scratch_buffer_size=scratch.bytes(),
rank=comm_group.my_rank,
)
for algorithm in algorithms:
if algorithm.name == "default_allreduce_nvls_zero_copy":
return comm_group, algorithm, scratch
pytest.skip("default_allreduce_nvls_zero_copy is not available")
def _torch_tensor_from_pool_buffer(torch, buffer, nelems: int):
return torch.utils.dlpack.from_dlpack(buffer.to_dlpack(data_type=str(torch.float32), shape=[nelems]))
def _run_nvls_zero_copy(algorithm, comm_group, buffer, stream) -> None:
ret = algorithm.execute(
comm=comm_group.communicator,
input_buffer=buffer.data(),
output_buffer=buffer.data(),
input_size=buffer.bytes(),
output_size=buffer.bytes(),
dtype=DataType.float32,
op=ReduceOp.SUM,
stream=stream.ptr,
nblocks=0,
nthreads_per_block=0,
symmetric_memory=True,
accum_dtype=DataType.float32,
)
assert ret == 0
@parametrize_mpi_groups(2, 4, 8)
def test_gpu_buffer_pool_allreduce_nvls_zero_copy_timing(mpi_group: MpiGroup):
torch = pytest.importorskip("torch")
if not torch.cuda.is_available():
pytest.skip("Torch CUDA is not available")
if not is_nvls_supported():
pytest.skip("NVLS is not supported")
if not _same_host(mpi_group.comm):
pytest.skip("NVLS zero-copy test requires all ranks on the same host")
torch.cuda.set_device(cp.cuda.Device().id)
comm_group, algorithm, scratch = _build_nvls_zero_algorithm(mpi_group)
stream = cp.cuda.Stream(non_blocking=True)
message_sizes = (256 * 1024, 1024 * 1024)
element_size = torch.empty((), dtype=torch.float32, device="cuda").element_size()
n_warmup = 3
n_iters = 10
pool = GpuBufferPool(sum(nbytes + 4096 for nbytes in message_sizes))
expected = float(comm_group.nranks * (comm_group.nranks + 1) // 2)
live_tensors = []
graphs = []
try:
for nbytes in message_sizes:
nelems = nbytes // element_size
buffer = pool.allocate(nbytes, alignment=4096)
tensor = _torch_tensor_from_pool_buffer(torch, buffer, nelems)
tensor.fill_(float(comm_group.my_rank + 1))
torch.cuda.synchronize()
mpi_group.comm.barrier()
_run_nvls_zero_copy(algorithm, comm_group, buffer, stream)
stream.synchronize()
assert torch.allclose(tensor, torch.full_like(tensor, expected))
tensor.fill_(float(comm_group.my_rank + 1))
torch.cuda.synchronize()
mpi_group.comm.barrier()
graph = capture_graph(stream, lambda: _run_nvls_zero_copy(algorithm, comm_group, buffer, stream))
graphs.append(graph)
graph.launch(stream)
stream.synchronize()
assert torch.allclose(tensor, torch.full_like(tensor, expected))
for _ in range(n_warmup):
graph.launch(stream)
stream.synchronize()
mpi_group.comm.barrier()
start = cp.cuda.Event()
end = cp.cuda.Event()
start.record(stream)
for _ in range(n_iters):
graph.launch(stream)
end.record(stream)
end.synchronize()
mpi_group.comm.barrier()
elapsed_us = cp.cuda.get_elapsed_time(start, end) * 1000.0 / n_iters
all_elapsed_us = mpi_group.comm.allgather(elapsed_us)
if comm_group.my_rank == 0:
avg_us = max(all_elapsed_us)
print(
f"default_allreduce_nvls_zero_copy graph with GpuBufferPool: "
f"nranks={comm_group.nranks}, nbytes={nbytes}, avg={avg_us:.2f} us"
)
live_tensors.append(tensor)
del buffer
finally:
for graph in graphs:
graph.close()
live_tensors.clear()
torch.cuda.synchronize()
AlgorithmCollectionBuilder.reset()
del scratch