Mooncake EP Design

Mooncake EP Design#

Mooncake EP is the expert-parallel communication runtime used for MoE token dispatch and combine. It follows the DeepEP low-latency programming model while adding Mooncake transport integration and rank activeness awareness.

This document explains the runtime at a level useful for developers maintaining Mooncake EP or integrating it into inference engines.

Goals#

Mooncake EP is designed to:

  • provide low-latency dispatch/combine operations for expert-parallel MoE inference;

  • keep the Python programming model close to DeepEP low-latency mode;

  • use Mooncake device transports for fast intra-node and inter-node movement;

  • detect failed source ranks through timeout-aware kernels;

  • interoperate with Mooncake PG for bootstrap metadata exchange and rank-health state.

High-level data flow#

MoE inference with Mooncake EP has three phases:

  1. Dispatch: each rank sends token hidden states to the ranks that own the selected experts. The receiver packs tokens by local expert.

  2. Expert compute: each rank runs its local experts over the packed inputs.

  3. Combine: expert outputs are routed back to the original token owners and reduced with routing weights.

        flowchart LR
    A[Local tokens x + topk_idx] --> B[dispatch]
    B --> C[Packed local expert inputs]
    C --> D[Local expert kernels]
    D --> E[Packed expert outputs]
    E --> F[combine]
    F --> G[Combined local token outputs]

Relationship with Mooncake Backend (PG)#

Mooncake EP is constructed from a torch.distributed process group:

from mooncake.mooncake_ep_buffer import Buffer

buffer = Buffer(group, num_ep_buffer_bytes)

The group should normally be a Mooncake Backend process group. EP uses it to exchange RDMA memory-region metadata, QP information, GID/LID information, and IPC handles. The backend active-rank state is also used by the fallback path and by peer synchronization.

After PG recovery changes membership, call Buffer.update_ep_member() to refresh EP peer metadata and QPs before relying on the recovered ranks for dispatch and combine.

Runtime objects#

Python wrapper#

mooncake.mooncake_ep_buffer.Buffer is the user-facing wrapper. It owns:

  • the process group;

  • the native EP buffer runtime;

  • the fallback flag;

  • Python fallback buffers used when fast path is unavailable.

Buffer.connect() exchanges peer metadata. It first tries the RDMA/IBGDA path, then exchanges IPC handles for intra-node P2P. If neither fast path is usable, it falls back to a Python implementation based on PyTorch collectives.

Native buffer#

The native MooncakeEpBuffer owns or references:

  • rank and world-size metadata;

  • a GDR workspace buffer;

  • P2P and RDMA device transports;

  • a communication stream;

  • temporary workspace used by dispatch/combine kernels.

If an external Transfer Engine is supplied by a higher layer, EP can reference device transports owned by that engine; otherwise EP creates and owns the transports itself.

Buffer layout#

EP allocates paired send/receive buffers for double-buffered dispatch/combine. Each pair contains:

  • RDMA send signal buffer;

  • RDMA receive signal buffer;

  • RDMA send data buffer;

  • RDMA receive data buffer.

The workspace size returned by Buffer.get_ep_buffer_size_hint() is derived from num_max_dispatch_tokens_per_rank, hidden, num_ranks, and num_experts. Size this for peak dispatch demand, not the average request.

Fast paths and fallback#

Mooncake EP has three broad execution modes:

Mode

Purpose

Notes

IBGDA / RDMA fast path

Inter-node GPU memory movement.

Requires RDMA-capable environment and successful metadata/QP setup.

P2P / IPC fast path

Intra-node peer access such as NVLink.

Requires peer accessibility and IPC handle exchange.

Python fallback

Functional fallback for unsupported environments.

Slower; useful for correctness and limited testing.

The runtime reports whether IBGDA is disabled and whether the fast path is usable through native helper methods surfaced in Python. The Python wrapper updates its fallback flag after metadata exchange.

Intel XPU#

On Intel XPU the native buffer is a SYCL implementation (mooncake-ep/src/xpu/) with the same Python-facing methods, so mooncake_ep_buffer.Buffer, its outputs and its handle are unchanged. The Python wrapper and the ep_py.cpp bindings are shared with CUDA, as is the host bookkeeping that pairs each combine with its dispatch (mooncake_ep_dispatch_slots.h); the kernels and transport are XPU-specific. It provides the P2P / IPC fast path only:

  • Peers share one device buffer through Level Zero IPC handles, exchanged over the process group like the CUDA IPC handles. Ranks must be on one host and in one PID namespace, their devices need peer access to each other, and each rank must be allowed to pidfd_getfd the others: Yama ptrace_scope 0 or 1, or 2 with CAP_SYS_PTRACE. Otherwise, and for multi-node groups, the Python fallback is used.

  • Data is pulled rather than pushed: each rank stages what it sends in its own half of the buffer and only writes small ready flags to peers. On PCIe a GPU writing to several peers at once is far slower than one reading from several.

  • A token that picks several experts on one rank crosses PCIe once: dispatch pulls it once and copies it to each of those experts locally. When dispatch() is given topk_weights, combine also sums each token’s rows on the expert side, so one row per rank travels back. The partial sum is rounded to BF16 once.

  • Kernels run on the caller’s torch.xpu stream. Zero-copy combine and the elastic buffer are not supported.

JIT kernel compilation#

On NVIDIA GPUs, Mooncake EP compiles its dispatch and combine kernels, including the elastic buffer kernels, at runtime. Parameters such as hidden, num_experts, and num_topk are compile-time template arguments of these kernels. Without JIT, the library can only contain a fixed set of instantiations (for example, hidden in {2048, 2560, 3072, 3584, 4096, 5120, 6144, 7168, 8192}), and any other value fails with an “Unsupported” error. With JIT, the first call with a new combination of values runs nvcc to build that instantiation and caches the cubin on disk; later calls and later runs reuse it.

JIT is currently not supported on other backends.

At runtime this needs nvcc, the kernel headers, and a writable cache directory. Installed wheels ship the kernel headers, so they usually need no setup. See JIT kernel compilation for path lookup, cache behavior, and environment variables.

Metadata exchange#

During Buffer.connect(), ranks exchange:

  • RDMA memory-region address and key;

  • local QP numbers;

  • LID and GID information;

  • subnet prefix and interface ID;

  • CUDA IPC handles for local peer access;

  • current active-rank mask from the backend.

The exchange uses dist.all_gather() and dist.all_to_all() on the process group. For this reason, the process group must already be initialized and healthy before the EP buffer is constructed or refreshed.

Dispatch internals#

dispatch() takes local token hidden states and selected expert IDs. It sends tokens to expert-owner ranks and packs received tokens into local-expert-major layout.

Important inputs:

  • x: [num_tokens, hidden] token hidden states;

  • topk_idx: [num_tokens, top_k] global expert IDs;

  • active_ranks: [num_ranks] int32 rank-health tensor;

  • num_max_dispatch_tokens_per_rank: receive capacity per source rank;

  • num_experts: global expert count, divisible by num_ranks;

  • timeout_us: failure-detection timeout.

Important outputs:

  • packed local-expert input tensor;

  • per-local-expert receive counts;

  • source/layout metadata handle used by combine();

  • event/hook synchronization helpers.

When use_fp8=True, dispatch returns packed FP8 data and FP32 scales. The local expert path must either consume that format directly or dequantize before expert compute.

Combine internals#

combine() sends local expert outputs back to token-owner ranks and applies routing weights. It consumes the handle produced by the matching dispatch() call.

For zero-copy combine, call get_next_combine_buffer(handle), write expert outputs into that buffer, then call combine(..., zero_copy=True) with the same handle. Do not reuse a handle across unrelated dispatch/combine pairs.

Rank activeness and timeout behavior#

Mooncake EP receives an active_ranks tensor in both dispatch and combine. The native kernels poll receive signals from source ranks. If timeout_us is not -1 and a source rank does not make progress before the timeout, the kernel can mark active_ranks[src_rank] = 0 and skip that source.

This EP-level tensor is rank-level and should have shape [num_ranks]. It is related to, but not automatically identical to, the backend-level active-rank mask passed through MooncakeBackendOptions. Integrations should propagate health updates consistently between scheduling logic, PG state, and EP buffers.

Stream synchronization model#

Mooncake EP operations return an EventOverlap object and, optionally, a hook:

  • If return_recv_hook=False, call event.current_stream_wait() before using the output tensors on the current stream.

  • If return_recv_hook=True, call the returned hook() at the chosen overlap point.

  • If async_finish=True, the wrapper records extra tensors in the event helper to keep lifetimes safe for asynchronous use and CUDA graph scenarios.

Always make synchronization explicit when composing EP with custom expert kernels, CUDA graphs, or application-level streams.

Recovery integration#

When a rank fails, multiple layers must be updated:

  1. PG active-rank state must stop collectives from waiting for the failed rank.

  2. Scheduler / MoE routing should stop assigning tokens to unavailable experts.

  3. Replacement ranks should join through the PG elastic protocol.

  4. EP buffers should refresh peer metadata with update_ep_member() after the process group activates recovered ranks.

EP timeout detection can mark a rank inactive in the EP-level tensor, but higher layers still need to coordinate recovery and routing decisions.

Performance considerations#

  • Prefer fast path operation with working RDMA/IBGDA or P2P peer access.

  • Size num_max_dispatch_tokens_per_rank for the worst expected per-rank token count to avoid overflow.

  • Use return_recv_hook=True or async_finish=True only when the application deliberately overlaps communication and compute.

  • use_fp8=True reduces dispatch bandwidth but requires FP8-aware expert code or explicit dequantization.

  • Repeatedly reconstructing EP buffers is expensive; refresh membership only when PG membership changes.

Developer test checklist#

For EP changes, run correctness across:

  • BF16 and FP8 dispatch;

  • zero-copy and non-zero-copy combine;

  • synchronous event wait and hook-based synchronization;

  • fallback and fast path when available;

  • failure simulation with finite timeout_us;

  • multi-rank topologies where num_experts % num_ranks == 0.

Useful entry points:

# EP grid correctness test
python python/tests/ep/test_ep_grid.py

# EP wrapper smoke test
python python/tests/ep/test_mooncake_ep.py

# Intel XPU correctness test (exact payloads, hooks, two batches, timeouts)
python python/tests/ep/test_mooncake_ep_xpu.py --num-ranks 4 --num-experts 256

# Host-only test of the dispatch bookkeeping shared by every backend
ctest --test-dir build -R ep_dispatch_slots_test

Adapt launch commands to the target environment and number of GPUs.

Contents