Writing

Push or Pull? Communication Design in DeepEP and Mega MoE

Yanjie Huang·

MoE / GPU Systems

The direction of a transfer follows who knows its final address. DeepEP and Mega MoE illustrate how that choice shapes buffering, execution order, and GPU resource use.

On this page

Every MoE layer follows the same broad sequence: dispatch tokens to their selected experts, run the expert matrix multiplications, and combine the outputs on the original GPUs. The communication direction inside that sequence is a design choice, not an intrinsic property of expert parallelism.

DeepSeek has two implementations of this:

  • DeepEP handles only the communication. Dispatch and combine are each one kernel; the matrix multiplications in between are left to DeepGEMM.
  • Mega MoE (inside DeepGEMM) puts dispatch, both matrix multiplications and combine into a single kernel.

DeepEP pushes tokens during dispatch; Mega MoE has the receiver pull them. Both push during combine. The asymmetry is revealing: communication direction follows when the final address becomes known and whether arrival order can advance computation.

The analysis uses the DeepGEMM and DeepEP source code and covers the SM100 implementation of Mega MoE. Examples use the DeepSeek V4.1-Flash shape: 384 routed experts, 6 picked per token, hidden size 5120, expert parallelism across 8 GPUs.

At a glance

  • Same job, two approaches. In DeepEP’s dispatch, the sender pushes tokens over and they are reordered on arrival. In Mega MoE, the receiver first exchanges counts, then pulls each token straight to where the matrix multiplication will read it. Both push in combine.
  • Why only dispatch flips: does the writer know the final address at the moment it writes? In combine, the address traveled with the token, so it pushes. In dispatch, the final position can only be computed once every GPU’s counts are in. DeepEP doesn’t wait, pushing first and reordering later; Mega MoE waits, then pulls each token straight into place.
  • Pulling turns scatter into gather. The pull walks the matrix in row order, so low blocks fill first and the GEMM can start while the pull is still going. This only pays off when communication and computation share one kernel, which is exactly why DeepEP doesn’t do it.
  • Mega MoE pays a little so data lands in place: an extra count exchange, about 1.35× the NVLink traffic, and NVLink-domain only. In return, data arrives exactly where the computation reads it.

1. Two execution boundaries for the same MoE layer

DeepEP + DeepGEMM runs several kernels one after another:

  1. Dispatch: each source GPU pushes its tokens into the receive buffer on the destination GPU. The buffer is split by source GPU, one region each. At the same time the GPUs exchange counts of how many tokens each expert receives.
  2. Reorder: the destination GPU rearranges what it received into a matrix with each expert’s tokens stored contiguously, and hands it to the GEMM.
  3. GEMM: DeepGEMM runs both matrix multiplications in a separate kernel.
  4. Combine: the GPU holding the expert pushes the results back to the source GPU, which sums them locally.

Mega MoE uses one kernel:

  1. Exchange counts: each GPU first tells the destination GPU who is coming and how many.
  2. Pull: the destination GPU pulls tokens from the source GPUs and places them directly where the first GEMM layer will read them.
  3. GEMM: both matrix multiplications run while the pull is still going on.
  4. Combine: when the second layer finishes, its results are pushed straight back to the source GPU, which sums them locally.
Dispatch: DeepEP pushes, Mega MoE pulls
Figure 1. Data flow in DeepEP and Mega MoE

The two have a lot in common:

  • Both run dispatch → two matrix multiplications → combine;
  • Both filter by expert before sending: a token goes only to the GPUs that hold its chosen experts, never to all of them;
  • Both read and write other GPUs’ memory directly through symmetric memory;
  • Both exchange counts first and then synchronize across GPUs;
  • Both push in combine.

One detail differs: the unit of sending. Several of a token’s 6 experts may sit on the same GPU. DeepEP deduplicates per GPU, sending one copy to each GPU and expanding it on arrival (source). Mega MoE does not deduplicate: every token–expert pair gets its own row, so two experts on the same GPU pull the same token twice (source). With uniform routing, DeepEP sends 4.43 copies per token on average and Mega MoE sends 6, so Mega MoE uses about 1.35× the NVLink traffic. As the following analysis shows, this is one of the prices Mega MoE pays so that data can be computed on as soon as it lands.

2. Address ownership determines push or pull

Only dispatch is reversed

DeepEP’s combine writes results back to the source GPU using the origin information recorded during dispatch (source). Mega MoE does the same (source).

The shared choice of push for combine rules out a general preference for pull. The relevant distinction is address ownership: can the writer determine the final destination using information it already has?

The test: when the write happens, does the writer know the final address?

In combine, the destination address is fixed by three numbers: the source GPU, the source token index, and which of the token’s choices this is. All three travelled with the token to the expert’s GPU during dispatch (source). Writing back only needs a local table lookup and depends on no other GPU. The writer knows the address, so it pushes.

In dispatch, the destination address is “row k of expert E’s matrix”, where k is the number of tokens, across all source GPUs, that also chose E and come before this one. Computing k requires every source GPU’s count for E. At the moment of writing, no GPU knows it.

That leaves two options:

  • DeepEP doesn’t wait. Data is pushed right away into a region split by source GPU, a position that can be computed from local information alone. Counts and data travel together, and the data is reordered locally after it arrives.
  • Mega MoE waits. It finishes exchanging counts first, lets the destination GPU compute k itself, and then pulls each token straight to its final position, with no reorder afterwards.

How Mega MoE “waits”

Before any data moves, Mega MoE writes two things to the destination GPU (source):

  • A list: one 4-byte integer per token–expert pair, recording “token i on source GPU r chose expert E”.
  • Lengths: how many tokens each source GPU sends to each expert, plus a signal that every GPU has finished writing.

Both are needed. Without the list, the destination does not know where to read from; without the lengths, it does not know where to stop, or whether the source has finished writing. After that, all GPUs synchronize once more before the pull begins.

This is the cost of pulling: the exchange must complete serially before any data transfer begins. According to comments in the source, the counting takes about 6.5 µs and writing the list about 2 µs (source).

Five constraints: why each side is right

Five constraints explain the different choices. They connect the location of the consumer, work available during synchronization, remote-memory access, buffer backpressure, and completion detection.

  1. Mega MoE includes the computation, while DeepEP handles only communication;
  2. Mega MoE also has a shared expert to compute, and can compute it while waiting for counts;
  3. Mega MoE is itself the user of its buffer, so it knows when space frees up and can stop and wait on its own; DeepEP cannot;
  4. The side that pulls naturally knows when all the data has arrived;
  5. Pulling needs direct access to other GPUs’ memory, which only works within one NVLink domain.

Mapped to the source:

ConstraintDeepEPMega MoE
① Where the dispatch consumer isIn another kernel; the kernel boundary is a global barrier, so arrival order has no value to it and it should only optimize “the whole transfer ends as early as possible” → don’t wait for counts; push first, reorder laterIn the same kernel and able to wait per block (source) → arrival order = compute order is what gives value → pull
② The gap while waiting for counts at the startNothing in the kernel to compute, so waiting is pure idle timeIssue the shared-expert L1 first (and as much L2 as fits in one wave), since it doesn’t depend on dispatch; the comment calls it the “dispatch gap” (source)
③ Can the peer’s memory be addressed directlyMust support cross-node RDMA/hybrid, so this cannot be assumedOnly within an NVLink domain (source; DeepInfer’s enabling condition: source)
④ Can the buffer be boundedThe receiver cannot make the sender stop → allocate for the worst case, [num GPUs, Tmax, record]The one waiting for a free slot is the receiver itself, reading its local l1_empty_count (source) → natural backpressure → a ring pool sized by the maximum number of blocks alive at once (source)
⑤ Who knows “the data has all arrived”Push: the receiver only knows through a completion signal from every sender, and the number of signals grows with EP size (MoK cites waiting on up to 71 peers)Pull: arrival is a local fact for the receiver (its local l1_full_count fills up, source). Note Mega MoE is not free of cross-GPU signals: they are concentrated in a single sync at the start (count exchange + NVLink barrier, source), after which every arrival check is local

Two notes:

  • On ②: the shared expert fills only part of this gap. The warps that post-process GEMM results (the epilogue) must first wait for dispatch to pass the first cross-GPU sync (source), and TMEM, which holds matrix-multiply results, has only 2 slots (source). So each SM can run at most 2 shared-expert tasks during the gap before it has to wait. The timeline in section 4 shows this.
  • On ⑤: this point comes from Cursor’s Mixture-of-Kittens. In a training setting (NVL72, EP64) they independently made the same choice, pull for dispatch and push for combine, for reasons that largely match ①, ④ and ⑤.

The table leads to one rule: the direction of communication is decided by who knows the final address and who needs to control the pace. When only the receiver can compute the final address and it wants data in the order it will compute, pull; when the sender already knows the final address, push. That is why both sides push in combine.

In addition, pulling and fusion depend on each other. Fusion is what makes the order of pulling matter; pulling in turn lets the receiver control the pace, so the buffer can be a small ring that is reused, which is what lets one kernel handle a large batch of tokens.

3. Arrange arrival order around the consumer

Pulling becomes useful when the receiver can request rows in the order its GEMM will consume them. The next question is therefore how a transfer completes whole compute tiles early, rather than merely moving individual tokens early.

Blocks exist only in the final matrix

The GEMM does not process an expert all at once. It splits the expert’s tokens into blocks of BLOCK_M rows and computes them block by block. Whether a block can start depends only on whether all of its BLOCK_M tokens have arrived.

On its way from the source GPU to the GEMM, a token passes through at most three layouts:

LayoutWhereRow orderHas blocks?
① Source input xSource GPU HBMSource token indexNo
② DeepEP receive bufferDestination GPU HBM[source GPU][slot], slots assigned by the source (source)No
③ Expert pool = L1’s A matrixDestination GPU HBMBy local expert, each expert padded to a multiple of BLOCK_MYes, one block per BLOCK_M rows

Mega MoE goes ① → ③; DeepEP goes ① → ② → ③.

A parallel loop finishes in the order of the index it walks

Mega MoE: the pull loop’s index is the row number of ③ (source). All pulling warps start at row 0 of ③ and move forward, so blocks with lower row numbers fill first. The scheduler also hands out GEMM tasks starting from the lowest blocks (source). Pull order and compute order match. Each time a token is pulled, that block’s counter goes up by one (source); when the loading warp sees the counter full, it starts computing that block (source).

DeepEP: on the source GPU, the loop index is the source token index (source); on the destination GPU, the reorder loop scans ② row by row and claims a position in ③ with one atomic add per row (source). Both steps walk the input in order and scatter into the output, so all blocks of ③ fill a little at a time, together, and almost none is complete until the end.

Who runs the loopLoop indexWhat finishes first
Pull = gatherReceiverOutput positionThe low end of the output
Push = scatterSenderInput positionThe low end of the input is sent first; the output fills sparsely everywhere

Gather has one requirement: you must be able to work out, from an output position, where the input is. The list and lengths in section 2 are the price paid for that.

For DeepEP, the order in which blocks fill does not matter. Its consumer, DeepGEMM, is another kernel that cannot start until DeepEP has finished entirely, so a block that arrives early cannot be used early. DeepEP only needs the whole transfer to finish as soon as possible, which is exactly constraint ①.

A small example

GPU 1 holds experts E0 and E1, and BLOCK_M = 2. Tokens on GPU 0: a picks E0, b picks E1, c picks E0. Tokens on GPU 1: d picks E0, f picks E1.

Row of ③012345
Tokena (GPU 0)d (GPU 1)c (GPU 0)emptyb (GPU 0)f (GPU 1)
BlockE0 block 0E0 block 0E0 block 1E0 block 1E1 block 0E1 block 0
  • Mega MoE pulls rows 0, 1, 2, 4, 5 in that order. E0’s block 0 fills first and can be computed right away.
  • DeepEP’s reorder scans ② in the order a, b, c, d, f, so E0’s block 0 is not full until it reaches d. And which token comes first inside E0 depends on which atomic add wins.

In addition, within an expert, tokens from different source GPUs take turns (a, d, c rather than a, c, d). This is Mega MoE’s round-robin rule (source). It doesn’t decide which block fills first; it decides which GPU each block reads from, so that reads at any moment are spread across all source GPUs instead of concentrating on one.

This gives a more general lesson: to let the consumer start as early as possible, organize the producer’s loop in the order the consumer needs. The same idea is behind output-stationary GEMM, where each thread block owns a fixed tile of the output.

4. Place communication within the SM budget

A receiver-driven transfer still needs GPU resources to issue loads, track progress, and feed the matrix pipeline. The remaining design question is how to divide those resources without leaving an expensive part of the GPU idle.

Three ways to divide the work

How work is dividedExampleWhere the idle time goes
By time: one kernel finishes, the next startsDeepEP dispatch → DeepGEMM → DeepEP combineTensor cores idle during communication; NVLink idle during compute
By SM: some SMs communicate, others computeWhen DeepEP’s combine runs alongside the shared expert, DeepEP takes only some of the SMs (DeepInfer caps it at min(theoretical, total − min(32, total//4)), source); Cursor’s MoK also sets aside dedicated communication SMsTensor cores on the communication SMs stay idle; the split is fixed and needs tuning
By warp, inside each SMMega MoE: every SM has 4 pull warps, 4 GEMM-related warps (two that load data, one that issues the matrix-multiply instructions, one scheduler) and 8 epilogue warpsSpread across every SM, and the roles that sit idle hold very few resources

Note that splitting work inside the SM does not mean an SM computes the tokens it pulled itself. When a pull warp finishes a token, it updates a global counter (source), and the scheduler decides dynamically which SM computes that block (source). Both approaches hand work over through global counters. The real difference is how resources are split and where the idle time goes.

Mega MoE’s resource budget

Registers: 48 per thread for the pull warps, 40 for the GEMM-related warps, 208 for the epilogue warps (source). That adds up to 128×48 + 128×40 + 256×208 = 64512, exactly filling the register file (source).

Shared memory: pulling takes about 22 KB, 2 KB of expert counts plus a 5 KB staging buffer for each of the 4 warps. Each extra GEMM pipeline stage needs about 19 KB (source, for BLOCK_M = 32), so pulling costs roughly one pipeline stage.

Put directly, Mega MoE leaves pulling with the bare minimum and gives most of each SM to the matrix multiplications and the epilogue. This works because pulling is done mostly by the TMA hardware and needs almost no computation of its own.

When each kind of warp is busy

When each warp role is busy inside the Mega MoE kernel
Figure 2. When each kind of warp is busy (8192 tokens per GPU, schematic)
  • Pull warps: exchange counts first (source), then pull in row order (source). When done, they stop and wait for all second-layer results to be written (source), then clean up the workspace and do the second cross-GPU sync, overlapping with combine (source).
  • Scheduler warp: schedules the shared expert’s first layer (source), waits for the counts (source), then schedules each expert’s first layer, and after that interleaves the two layers.
  • Load and MMA warps: work whenever there is a task; once scheduling ends, they sit idle through the whole combine phase.
  • Epilogue warps: blocked at the start (source), then process the results of both layers, and finally turn into combine warps that do the summing (source onward).

Every kind of warp has a stretch of idle time: the pull warps near the end of compute, the GEMM-related warps during combine, the epilogue warps at the start. Mega MoE can accept this because the roles that sit idle hold very few resources, while the most resource-hungry role, the epilogue warps, stays busy from the start of compute almost to the end. The same warp also changes jobs over time: epilogue warps later do combine, and pull warps later do cleanup.

This also explains why Mega MoE doesn’t set aside a group of SMs just for pulling: then the idle part would be whole SMs, tensor cores and all their registers included.

5. What fusion changes

Mega MoE spends extra metadata work, some additional NVLink traffic, and a portion of each SM’s resources to make data arrive directly in compute order. Fusion makes that ordering valuable: a complete block can begin computation before dispatch as a whole has finished.

  • Pull instead of push: an extra serial count exchange, in return for data arriving in compute order;
  • No per-GPU deduplication: about 1.35× the NVLink traffic, in return for computing as soon as data lands, with no second expansion;
  • Splitting work by warp inside each SM: every SM gives up a little, in return for every tensor core doing work;
  • A reused ring buffer: the receiver has to control the pace, so it only works within an NVLink domain where peer memory is directly addressable.

DeepEP serves a different boundary. Its GEMM consumer is a separate kernel, and its communication layer must support cross-node paths. Early block arrival cannot cross that kernel boundary, so minimizing the completion time of the whole transfer is the appropriate target.

Several questions remain open: why the epilogue warps are blocked at the start (the source comment only says it must not run alongside the barrier before pulling); whether the serial exchange at the start could be made cheaper; and how much each of these differences matters under decode and under prefill. The last question will be analyzed in a follow-up post, starting from the number of tokens each expert receives.

References