microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-15 11:44:56 +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

31
test/python/ext/ep/test_internode_multirank.py
View File

@@ -85,8 +85,11 @@ def inplace_unique(x: torch.Tensor, num_slots: int):
def main():
rank, num_ranks, local_rank, group = init_dist()
from mscclpp import CommGroup
from mscclpp.ext import ep
ep_group = CommGroup(torch_group=group)
NUM_MAX_NVL_PEERS = _detect_local_world_size()
assert (
num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS
@@ -139,7 +142,7 @@ def main():
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank)
# Buffer config for internode HT: needs num_rdma_bytes > 0. Size buffers
# Runtime config for internode HT: needs num_rdma_bytes > 0. Size buffers
# using max(hidden, bench_hidden) so the optional bench phase fits.
cfg = ep.Config(
int(os.environ.get("MSCCLPP_EP_NSM", "152")),
@@ -160,16 +163,18 @@ def main():
flush=True,
)
print(f"[rank {rank}] creating Buffer", flush=True)
buf = ep.Buffer(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False)
print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True)
buf = ep.ExpertParallelRuntime(
ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=False
)
print(
f"[rank {rank}] Buffer created is_available={buf.is_available()} "
f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()} "
f"is_internode={buf.is_internode_available()}",
flush=True,
)
assert buf.is_available() and buf.is_internode_available()
ref_rank, ref_rdma_rank, ref_exp, ref_in_rank, _ = buf.runtime.get_dispatch_layout(
ref_rank, ref_rdma_rank, ref_exp, ref_in_rank, _ = buf.get_dispatch_layout(
topk_idx, num_experts, None, False, False
)
assert torch.allclose(ref_rank, num_tokens_per_rank)
@@ -203,7 +208,7 @@ def main():
send_rdma_head,
send_nvl_head,
_event,
) = buf.runtime.internode_dispatch(
) = buf.internode_dispatch(
x,
None,
topk_idx,
@@ -226,7 +231,7 @@ def main():
)
dist.barrier(group=group)
_skip_verify = os.environ.get("MSCCLPP_EP_SKIP_VERIFY","0") in ("1","true","True")
_skip_verify = os.environ.get("MSCCLPP_EP_SKIP_VERIFY", "0") in ("1", "true", "True")
# Validate recv buffer: for each source rank i, the block carries value i.
assert recv_x.dim() == 2 and recv_x.size(1) == hidden
start = 0
@@ -248,7 +253,7 @@ def main():
# the various *_channel_prefix_matrix tensors can still be in flight on
# the comm stream when combine launches, producing a deadlock inside the
# combine forwarder (NVL check never advances). Investigate proper
# stream-dependency hand-off in Buffer::internode_dispatch.
# stream-dependency hand-off in ExpertParallelRuntime.internode_dispatch.
torch.cuda.synchronize()
dist.barrier(group=group)
@@ -261,7 +266,7 @@ def main():
# matrices passed here must be the RECEIVER-side ones returned by dispatch
# (`recv_rdma_channel_prefix_matrix`, `recv_rdma_rank_prefix_sum`,
# `recv_gbl_channel_prefix_matrix`) — not the sender-side ones.
combined_x, combined_topk_weights, _ = buf.runtime.internode_combine(
combined_x, combined_topk_weights, _ = buf.internode_combine(
recv_x,
recv_topk_weights,
recv_src_meta,
@@ -319,7 +324,7 @@ def main():
print(f"[bench] skip: topk={bench_num_topk} > experts={bench_num_experts}", flush=True)
return
# Respect the Buffer's pre-sized num_nvl_bytes / num_rdma_bytes budget.
# Respect the runtime's pre-sized num_nvl_bytes / num_rdma_bytes budget.
per_peer_nvl = num_nvl_bytes // max(1, num_ranks)
per_peer_rdma = num_rdma_bytes // max(1, num_ranks)
if bench_hidden * x.element_size() > min(per_peer_nvl, per_peer_rdma):
@@ -327,7 +332,7 @@ def main():
print(
f"[bench] skip: hidden={bench_hidden} bytes/row={bench_hidden * x.element_size()} "
f">= min(per-peer NVL {per_peer_nvl}, RDMA {per_peer_rdma}). "
f"Rerun with a larger Buffer or smaller hidden.",
f"Rerun with a larger runtime or smaller hidden.",
flush=True,
)
return
@@ -362,7 +367,7 @@ def main():
x_b = torch.ones((bench_tokens, bench_hidden), dtype=torch.bfloat16, device="cuda") * float(rank)
def _dispatch():
return buf.runtime.internode_dispatch(
return buf.internode_dispatch(
x_b,
None,
topk_idx_b,
@@ -386,7 +391,7 @@ def main():
def _combine(dout):
rx, _rxs, _rti, rtw, _lst, _rpm, _gpm, rrcpm, rrps, rgpm, _rgps, rsm, sh_rdma, sh_nvl, _ev = dout
buf.runtime.internode_combine(
buf.internode_combine(
rx,
rtw,
rsm,

29
test/python/ext/ep/test_intranode_multirank.py
View File

@@ -5,7 +5,7 @@
Launch with:
torchrun --nproc_per_node=<N> test/python/ext/ep/test_intranode_multirank.py
Tests that Buffer::sync() succeeds across N GPUs on a single node and that
Tests that ExpertParallelRuntime sync succeeds across N GPUs on a single node and that
a round-trip dispatch + combine preserves data (sum of top-k weighted copies).
Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass
@@ -65,8 +65,11 @@ def inplace_unique(x: torch.Tensor, num_slots: int):
def main():
rank, num_ranks, local_rank, group = init_dist()
from mscclpp import CommGroup
from mscclpp.ext import ep
ep_group = CommGroup(torch_group=group)
# Small settings for functional check
num_tokens = 128
hidden = 1024
@@ -104,7 +107,7 @@ def main():
# Token payload = rank id (cast to bf16) so we can check correctness
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank)
# Allocate Buffer (intranode only: num_rdma_bytes=0). Size the NVL buffer
# Allocate runtime (intranode only: num_rdma_bytes=0). Size the NVL buffer
# using max(hidden, bench_hidden) so the optional bench phase fits.
cfg = ep.Config(
int(os.environ.get("MSCCLPP_EP_NUM_SMS", "20")),
@@ -121,13 +124,13 @@ def main():
flush=True,
)
print(f"[rank {rank}] creating Buffer", flush=True)
buf = ep.Buffer(group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False)
print(f"[rank {rank}] Buffer created is_available={buf.is_available()}", flush=True)
print(f"[rank {rank}] creating ExpertParallelRuntime", flush=True)
buf = ep.ExpertParallelRuntime(ep_group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=0, low_latency_mode=False)
print(f"[rank {rank}] ExpertParallelRuntime created is_available={buf.is_available()}", flush=True)
assert buf.is_available()
# get_dispatch_layout sanity
ref_rank, _, ref_exp, ref_in_rank, _ = buf.runtime.get_dispatch_layout(topk_idx, num_experts, None, False, False)
ref_rank, _, ref_exp, ref_in_rank, _ = buf.get_dispatch_layout(topk_idx, num_experts, None, False, False)
assert torch.allclose(ref_rank, num_tokens_per_rank)
assert torch.allclose(ref_exp, num_tokens_per_expert)
assert torch.allclose(ref_in_rank, is_token_in_rank)
@@ -147,7 +150,7 @@ def main():
recv_src_idx,
send_head,
_event,
) = buf.runtime.intranode_dispatch(
) = buf.intranode_dispatch(
x,
None,
topk_idx,
@@ -188,7 +191,7 @@ def main():
handle_rank_prefix_matrix = rank_prefix_matrix
handle_channel_prefix_matrix = recv_channel_prefix_matrix
combined_x, combined_topk_weights, _ = buf.runtime.intranode_combine(
combined_x, combined_topk_weights, _ = buf.intranode_combine(
recv_x,
recv_topk_weights,
handle_recv_src_idx,
@@ -246,14 +249,14 @@ def main():
return
# Rebuild inputs at bench size. Keep same layout recipe as above but at
# larger (num_tokens, hidden); Buffer is sized off the original cfg+hidden,
# larger (num_tokens, hidden); runtime is sized off the original cfg+hidden,
# so bench must fit within num_nvl_bytes. If it doesn't, we skip.
if bench_hidden * x.element_size() > (num_nvl_bytes // max(1, num_ranks)):
if rank == 0:
print(
f"[bench] skip: hidden={bench_hidden} bytes/row={bench_hidden * x.element_size()} "
f"> per-peer budget {num_nvl_bytes // num_ranks}. "
f"Rerun with a larger Buffer or smaller hidden.",
f"Rerun with a larger runtime or smaller hidden.",
flush=True,
)
return
@@ -281,7 +284,7 @@ def main():
x_b = torch.ones((bench_tokens, bench_hidden), dtype=torch.bfloat16, device="cuda") * float(rank)
def _dispatch():
return buf.runtime.intranode_dispatch(
return buf.intranode_dispatch(
x_b,
None,
topk_idx_b,
@@ -315,7 +318,7 @@ def main():
# topk_idx/topk_weights (those require num_experts > 0). We still get
# send_head/rank_prefix_matrix/channel_prefix_matrix/recv_src_idx out
# of dispatch -- enough to drive combine.
return buf.runtime.intranode_dispatch(
return buf.intranode_dispatch(
x_b,
None,
None,
@@ -335,7 +338,7 @@ def main():
def _combine(dout):
rx, _rxs, _rti, rtw, _lst, rpm, _cpm, rcpm, rsi, sh, _ev = dout
buf.runtime.intranode_combine(
buf.intranode_combine(
rx,
rtw,
rsi,

199
test/python/ext/ep/test_low_latency_multirank.py
View File

@@ -3,7 +3,8 @@
"""Multi-rank low-latency functional test for mscclpp_ep.
Launch with (intra-node, 8 GPUs):
torchrun --nproc_per_node=8 test/python/ext/ep/test_low_latency_multirank.py
torchrun --nproc_per_node=8 test/python/ext/ep/test_low_latency_multirank.py \
--num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256
Launch with (2 nodes, 1 GPU per node -- DeepEP's recommended LL topology):
# node 0:
@@ -33,9 +34,9 @@ we need for an LL port smoke test. BF16-only (no FP8 check).
from __future__ import annotations
import argparse
import os
import random
import sys
# Disable ProcessGroupNCCL's HeartbeatMonitor before importing torch.distributed.
# It runs in a background thread polling the TCPStore; under mpirun, rank 0
@@ -47,6 +48,19 @@ import torch
import torch.distributed as dist
def parse_args():
parser = argparse.ArgumentParser(description="MSCCL++ EP low-latency multi-rank correctness/benchmark test")
parser.add_argument("--num-tokens", type=int, default=128)
parser.add_argument("--hidden", type=int, default=7168, help="LL kernels are compiled for a fixed hidden set")
parser.add_argument("--num-topk", type=int, default=8)
parser.add_argument("--num-experts", type=int, default=256)
parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness")
parser.add_argument("--bench-warmup", type=int, default=5)
parser.add_argument("--bench-iters", type=int, default=20)
parser.add_argument("--local-rank", "--local_rank", type=int, default=None, help=argparse.SUPPRESS)
return parser.parse_args()
def init_dist():
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
@@ -62,19 +76,21 @@ def init_dist():
def main():
args = parse_args()
rank, num_ranks, local_rank, group = init_dist()
from mscclpp import CommGroup
from mscclpp.ext import ep
ep_group = CommGroup(torch_group=group)
# Shrink the "bf16 precision" anchor to keep values small.
rank_offset = 128
assert num_ranks - rank_offset < 257, "too many ranks for bf16 precision anchor"
num_tokens = int(os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "128"))
hidden = int(
os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168")
) # LL kernels are compiled for a fixed set; see SWITCH_HIDDEN
num_topk = int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", "8"))
num_experts = int(os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", "256"))
num_tokens = args.num_tokens
hidden = args.hidden
num_topk = args.num_topk
num_experts = args.num_experts
assert num_experts % num_ranks == 0
num_local_experts = num_experts // num_ranks
@@ -93,63 +109,50 @@ def main():
for _ in range(min(10, num_tokens)):
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = -1
num_rdma_bytes = ep.Buffer.get_low_latency_rdma_size_hint(num_tokens, hidden, num_ranks, num_experts)
moe_comm = ep.MoECommunicator(
comm=ep_group,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden,
topk=num_topk,
max_tokens_per_rank=num_tokens,
mode=ep.MoEMode.LOW_LATENCY,
num_rdma_qps_per_rank=max(1, num_experts // num_ranks),
)
if rank == 0:
print(
f"[cfg] num_ranks={num_ranks} num_tokens={num_tokens} hidden={hidden} "
f"num_experts={num_experts} num_topk={num_topk} "
f"num_rdma_bytes={num_rdma_bytes}",
f"num_experts={num_experts} num_topk={num_topk}",
flush=True,
)
buf = ep.Buffer(
group,
num_nvl_bytes=0,
num_rdma_bytes=num_rdma_bytes,
low_latency_mode=True,
num_qps_per_rank=max(1, num_experts // num_ranks),
)
print(
f"[rank {rank}] Buffer created is_available={buf.is_available()} "
f"is_internode={buf.is_internode_available()}",
f"[rank {rank}] MoECommunicator created is_available={moe_comm.is_available()} "
f"is_internode={moe_comm.is_internode_available()}",
flush=True,
)
assert buf.is_available()
assert moe_comm.is_available()
dist.barrier(group=group)
torch.cuda.synchronize()
print(f"[rank {rank}] pre-dispatch", flush=True)
# --- Dispatch ---
# Return tuple (7 items):
# packed_recv_x, packed_recv_x_scales (optional, FP8-only),
# packed_recv_count, packed_recv_src_info, packed_recv_layout_range,
# event, hook
(
packed_recv_x,
_packed_recv_x_scales,
packed_recv_count,
packed_recv_src_info,
packed_recv_layout_range,
_event,
recv_hook,
) = buf.low_latency_dispatch(
dispatch_output_buffer = torch.empty(
(num_local_experts, num_ranks * num_tokens, hidden),
dtype=torch.bfloat16,
device="cuda",
)
dispatch_out, handle = moe_comm.dispatch(
x,
topk_idx,
num_tokens,
num_experts,
False,
False,
True, # use_fp8, async, return_recv_hook
topk_weights,
output_buffer=dispatch_output_buffer,
)
# Send phase launched on compute_stream; wait for local launch.
torch.cuda.synchronize()
dist.barrier(group=group)
print(f"[rank {rank}] dispatch-send done, calling hook", flush=True)
recv_hook() # Recv phase.
packed_recv_x = dispatch_out.tokens
packed_recv_count = dispatch_out.num_tokens_per_expert
packed_recv_layout_range = handle.layout_range
torch.cuda.synchronize()
print(f"[rank {rank}] post-dispatch", flush=True)
handle = (packed_recv_src_info, packed_recv_layout_range)
# packed_recv_x: [num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden]
# packed_recv_count: [num_local_experts] int32
@@ -162,7 +165,7 @@ def main():
expert_id = rank * num_local_experts + i
recv_count = int(packed_recv_count[i].item())
expected_count = int((all_topk_idx == expert_id).sum().item())
recv_layout_range = handle[1][i]
recv_layout_range = packed_recv_layout_range[i]
layout_sum = int((recv_layout_range & int_mask).sum().item())
assert (
recv_count == expected_count
@@ -187,30 +190,17 @@ def main():
# returns sum(x * weight) across experts.
simulated_gemm_x = packed_recv_x.clone()
out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
# Signature: (x, topk_idx, topk_weights, src_info, layout_range,
# num_max_dispatch_tokens_per_rank, num_experts,
# zero_copy, async, return_recv_hook, out)
src_info, layout_range = handle[0], handle[1]
combined_x, _event, _hook = buf.low_latency_combine(
simulated_gemm_x,
topk_idx,
topk_weights,
src_info,
layout_range,
num_tokens,
num_experts,
False,
False,
False, # zero_copy, async, return_recv_hook
out,
)
combined_x = moe_comm.combine(simulated_gemm_x, handle, out=out)
# Analytical expected: each token i, weighted sum over topk entries that
# are not -1. Every expert returns the original x[i] (since simulated
# gemm is identity), so the combine output should be
# x[i] * sum(topk_weights[i, j] for j where topk_idx[i,j] != -1).
weight_sum = topk_weights.masked_fill(topk_idx == -1, 0.0).sum(dim=1).view(-1, 1)
expected = (x.float() * weight_sum).to(torch.bfloat16)
# are not -1. Accumulate in the same top-k order as the kernel; multiplying
# by the pre-summed weights can differ by one BF16 ULP for large token IDs.
expected_f = torch.zeros_like(x, dtype=torch.float32)
x_f = x.float()
for j in range(num_topk):
weight_j = topk_weights[:, j].masked_fill(topk_idx[:, j] == -1, 0.0).view(-1, 1)
expected_f += x_f * weight_j
expected = expected_f.to(torch.bfloat16)
diff = (combined_x.float() - expected.float()).abs().max().item()
max_exp = expected.float().abs().max().item()
print(
@@ -225,56 +215,23 @@ def main():
print("PASS", flush=True)
# ------------------------------------------------------------------
# Optional benchmark (enable with MSCCLPP_EP_BENCH=1). Times dispatch
# and combine separately, reporting per-iter latency (max across ranks)
# and aggregate effective bandwidth (sum across ranks).
# Optional benchmark. Times dispatch and combine separately, reporting
# per-iter latency (max across ranks) and aggregate effective bandwidth
# (sum across ranks).
# ------------------------------------------------------------------
if os.environ.get("MSCCLPP_EP_BENCH", "0") != "1":
if not args.bench:
return
warmup = int(os.environ.get("MSCCLPP_EP_BENCH_WARMUP", "5"))
iters = int(os.environ.get("MSCCLPP_EP_BENCH_ITERS", "20"))
# Hoist dispatch's output tensors out of the timed loop. The largest
# (`packed_recv_x`, ~58 MB at 7K hidden) costs ~10us cumulative across
# the four torch::empty calls per iter; reusing them brings the bench
# in line with NCCL-EP `ep_bench` which preallocates output buffers.
num_local_experts = num_experts // num_ranks
bench_packed_recv_x = torch.empty(
(num_local_experts, num_ranks * num_tokens, hidden),
dtype=torch.bfloat16,
device="cuda",
)
bench_packed_recv_src_info = torch.empty(
(num_local_experts, num_ranks * num_tokens),
dtype=torch.int32,
device="cuda",
)
bench_packed_recv_layout_range = torch.empty(
(num_local_experts, num_ranks),
dtype=torch.int64,
device="cuda",
)
bench_packed_recv_count = torch.empty(
(num_local_experts,),
dtype=torch.int32,
device="cuda",
)
warmup = args.bench_warmup
iters = args.bench_iters
bench_dispatch_output_buffer = torch.empty_like(dispatch_output_buffer)
def _dispatch():
return buf.low_latency_dispatch(
return moe_comm.dispatch(
x,
topk_idx,
num_tokens,
num_experts,
False,
False,
False, # use_fp8, async, return_recv_hook
bench_packed_recv_x,
None, # x_scales (FP8 only)
bench_packed_recv_src_info,
bench_packed_recv_layout_range,
bench_packed_recv_count,
topk_weights,
output_buffer=bench_dispatch_output_buffer,
)
# Hoist combine's output-tensor allocation out of the timed loop so the
@@ -284,20 +241,8 @@ def main():
bench_out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
def _combine(dout, out_):
recv_x, _scales, _cnt, src_info_, layout_range_, _ev, _hk = dout
buf.low_latency_combine(
recv_x,
topk_idx,
topk_weights,
src_info_,
layout_range_,
num_tokens,
num_experts,
False,
False,
False,
out_,
)
dispatch_out_, handle_ = dout
moe_comm.combine(dispatch_out_.tokens, handle_, out=out_)
for _ in range(warmup):
_combine(_dispatch(), bench_out)
@@ -313,7 +258,7 @@ def main():
end_ev.record()
torch.cuda.synchronize()
disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters
recv_tokens = int(dout[2].sum().item()) # packed_recv_count summed over local experts
recv_tokens = int(dout[0].num_tokens_per_expert.sum().item())
dist.barrier(group=group)
start_ev.record()