# How Mega MoE Overlaps Expert Parallelism Communication with Computation in DeepGEMM

> Learn how Mega MoE overlaps expert parallelism communication with computation. Discover techniques that hide data movement using symmetric memory buffers and asynchronous NCCL streams.

- Repository: [DeepSeek/DeepGEMM](https://github.com/deepseek-ai/DeepGEMM)
- Tags: internals
- Published: 2026-04-19

---

**Mega MoE hides expert-parallel data movement behind FP8/FP4 GEMM kernels by using a symmetric memory buffer, asynchronous NCCL streams, and block-aligned token counts to overlap communication with computation.**

DeepGEMM, DeepSeek's open-source library for MoE optimizations, implements Mega MoE to eliminate the latency bottleneck inherent in expert parallelism. By overlapping the dispatch and combine communication phases with the heavy matrix multiplication workloads, the framework ensures GPUs remain fully utilized during token routing.

## Three Mechanisms for Communication-Computation Overlap

### Symmetric Memory Buffer for Zero-Copy Data Sharing

The foundation of the overlap strategy is a **symmetric memory buffer** that provides a shared view of inputs, top-k indices, and intermediate activations across all ranks. In [`deep_gemm/mega/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/mega/__init__.py) at lines 31-48, the buffer is allocated using `torch.distributed._symmetric_memory.empty` and rendezvoused across ranks via `symm_mem.rendezvous`. This mechanism allows all ranks to read and write the same logical regions without extra host-to-device copies or intermediate staging buffers.

### Asynchronous NCCL Communication on Dedicated Streams

The fused kernel `fp8_fp4_mega_moe`—defined in [`deep_gemm/mega/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/mega/__init__.py) at lines 10-18—launches non-blocking NCCL collectives on a **dedicated communication stream** while scheduling FP8/FP4 GEMM operations on the default compute stream. Because both streams attach to the same CUDA device, the GPU progresses the dispatch (`AllToAll` scatter) and combine (`AllToAll` gather) operations concurrently with the matrix multiplications. The actual overlap logic resides in the low-level C++/CUDA extension `_C.fp8_fp4_mega_moe`.

### Block-M Alignment for Irregular Token Counts

To prevent GEMM launch stalls caused by irregular token distribution across experts, Mega MoE enforces **Block-M alignment**. Token counts are padded to the GEMM block size using `_C.get_block_m_for_mega_moe`, as seen in [`deep_gemm/mega/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/mega/__init__.py) at lines 64-68. This alignment ensures the compute kernel launches with a fixed shape regardless of the exact number of tokens each rank receives, allowing the communication stream to continuously feed data into the already-running kernel without synchronization points.

## End-to-End Workflow

The complete overlap path follows a strict initialization and execution sequence:

1. **Initialize Distributed Group** – Each process joins a `torch.distributed.ProcessGroup` using `init_dist` from `deep_gemm.utils.dist`.
2. **Create Symmetric Buffer** – Call `deep_gemm.get_symm_buffer_for_mega_moe` to compute the required size (including space for raw inputs, scaled-factor tensors, and top-k metadata) and return a `SymmBuffer` containing CUDA views (`x`, `x_sf`, `topk_idx`).
3. **Copy Inputs into Buffer** – Copy FP8-packed activations and top-k indices into the shared slices (`buffer.x[:num_tokens].copy_(x_fp8)`, etc.) without extra host-to-device transfers.
4. **Launch Fused Kernel** – Execute `deep_gemm.fp8_fp4_mega_moe`, which internally issues the dispatch NCCL `AllToAll`, launches three FP8/FP4 GEMM kernels (L1 left, L1 right, L2), applies Swiglu activation in-place, and performs the combine NCCL `AllToAll`.
5. **Barrier and Cleanup** – Call `sym_buffer.group.barrier()` to guarantee all ranks complete both communication phases before destroying the buffer.

## Code Example: Fused Overlap Path

The following minimal snippet demonstrates the complete overlap path, contrasting with the non-overlapped baseline that uses separate dispatch and combine calls.

```python
import torch
import torch.distributed as dist
import deep_gemm
from deep_gemm.utils.dist import init_dist

# 1️⃣  Distributed init

rank, world, grp = init_dist(local_rank=0, num_local_ranks=8)

# 2️⃣  Create symmetric buffer

sym_buf = deep_gemm.get_symm_buffer_for_mega_moe(
    group=grp,
    num_experts=384,
    num_max_tokens_per_rank=8192,
    num_topk=6,
    hidden=7168,
    intermediate_hidden=3072,
    use_fp8_dispatch=True,
    activation='swiglu',
)

# 3️⃣  Prepare inputs (FP8‑packed activations + top‑k)

num_tokens = 4096
x = torch.randn((num_tokens, 7168), dtype=torch.bfloat16, device='cuda')
x_fp8, x_fp8_sf = deep_gemm.utils.per_token_cast_to_fp8(
    x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True
)
topk_weights, topk_idx = torch.topk(
    torch.randn(num_tokens, 384, device='cuda'), k=6
)

# 4️⃣  Copy into shared buffer (no extra host‑to‑device copies)

sym_buf.x[:num_tokens].copy_(x_fp8)
sym_buf.x_sf[:num_tokens].copy_(x_fp8_sf)
sym_buf.topk_idx[:num_tokens].copy_(topk_idx)
sym_buf.topk_weights[:num_tokens].copy_(topk_weights)

# 5️⃣  Transform weights (interleave + SF transpose) for L1 and L2

l1_w = torch.randn(384, 7168, 3072, dtype=torch.bfloat16, device='cuda')
l2_w = torch.randn(384, 3072, 7168, dtype=torch.bfloat16, device='cuda')
l1_fp4, l2_fp4 = deep_gemm.transform_weights_for_mega_moe(l1_w, l2_w)

# 6️⃣  Run fused kernel – communication & GEMM overlap internally

y = torch.empty((num_tokens, 7168), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(
    y,
    l1_fp4,
    l2_fp4,
    sym_buf,
    activation_clamp=10.0,
    fast_math=True,
)

# 7️⃣  Clean up

sym_buf.destroy()
dist.barrier()
dist.destroy_process_group()

```

## Comparison with Non-Overlapped Baseline

The baseline implementation—which serves as the performance reference in [`tests/test_mega_moe.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_mega_moe.py)—executes communication and computation sequentially. It first calls `deep_ep.ElasticBuffer.dispatch` to scatter tokens, performs the three GEMM operations (L1 left, L1 right, L2), applies the Swiglu activation, and finally calls `ep_buffer.combine` to gather results. This path incurs two full NCCL barriers and leaves the compute units idle during communication phases. In contrast, the fused `fp8_fp4_mega_moe` path overlaps both the dispatch and combine operations with the GEMM kernels, eliminating idle GPU time and achieving significantly higher effective TFLOPS as reported by the test harness.

## Summary

- **Symmetric memory buffers** created via `torch.distributed._symmetric_memory.empty` enable zero-copy data sharing across ranks, eliminating intermediate staging buffers.
- **Asynchronous NCCL streams** allow the `fp8_fp4_mega_moe` kernel to run dispatch and combine collectives concurrently with FP8/FP4 GEMM operations on separate CUDA streams.
- **Block-M alignment** via `_C.get_block_m_for_mega_moe` pads token counts to fixed GEMM block sizes, preventing kernel launch stalls from irregular expert loads.
- The fused implementation in [`deep_gemm/mega/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/mega/__init__.py) eliminates the two NCCL barriers required by the baseline, achieving higher effective TFLOPS and NVLink bandwidth utilization as measured in [`tests/test_mega_moe.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_mega_moe.py).

## Frequently Asked Questions

### What is expert parallelism in MoE models?

Expert parallelism is a distributed training and inference strategy where different experts in a Mixture-of-Experts (MoE) layer are placed on separate GPU ranks. Each rank holds a subset of experts and processes only the tokens routed to those experts, requiring scatter (dispatch) and gather (combine) communication phases to move tokens between ranks.

### How does the symmetric memory buffer avoid extra copies?

The symmetric memory buffer uses `torch.distributed._symmetric_memory.empty` to allocate a single CUDA memory region that is rendezvoused across all ranks via `symm_mem.rendezvous`. This creates a shared memory view where each rank can directly read and write the same logical tensor regions, eliminating the need for separate source buffers, intermediate staging areas, or explicit device-to-device copies between ranks.

### Why is Block-M alignment necessary for overlapping?

Block-M alignment pads the number of tokens per rank to the GEMM block size (retrieved via `_C.get_block_m_for_mega_moe`). Without this alignment, variable token counts across experts would cause the GEMM kernel to launch with irregular shapes, potentially stalling the compute stream and breaking the overlap with the asynchronous NCCL communication running on the dedicated stream.

### Where can I find the implementation details?

The primary implementation resides in [`deep_gemm/mega/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/mega/__init__.py), which defines the `SymmBuffer` class, `get_symm_buffer_for_mega_moe`, and the fused `fp8_fp4_mega_moe` entry point. The low-level C++/CUDA kernel that manages the NCCL streams and GEMM overlap is located in the `deep_gemm._C` extension module. End-to-end benchmarks and overlap metrics are available in [`tests/test_mega_moe.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_mega_moe.py).