Between Push and Pull: A Brief Analysis of the Communication Schemes in DeepGEMM Mega MoE and DeepEP

Chinese version

Expert parallelism (EP) for MoE does three things in every layer: send each token to the GPUs that hold the experts it picked (dispatch), run two matrix multiplications, and send the results back to the token’s own GPU to be summed (combine).

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.

One difference between the two is easy to miss: in DeepEP’s dispatch, the sender pushes the data over; in Mega MoE, the receiver pulls it. Yet both of them push in combine. This post starts from that difference, compares what the two implementations share and where they differ, and asks why they differ.

This post is based on 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.

TL;DR

• 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. What the two have in common

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.

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 we will see, this is one of the prices Mega MoE pays so that data can be computed on as soon as it lands.

2. Difference one: push versus 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).

If pulling were always better, combine should pull too; if pushing were always better, dispatch should not have changed. So the answer must lie in some difference between the two transfers.

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

Neither design is better in general; each is sensible under its own constraints, which can be summarized in five points:

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:

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. Difference two: from scatter to gather

Constraint ① says pulling is only worth it when arrival order equals compute order. So how does pulling make data arrive in compute order?

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:

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.

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.

• 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. Difference three: how the SMs divide the work

Once the receiver is the one pulling, the next question is which resources on the GPU it uses to pull.

Three ways to divide the work

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.

In other words, 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

• 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. Summary

In summary, almost every difference between Mega MoE and DeepEP is a facet of the same choice: pay a little elsewhere so that data lands exactly where the computation needs it.

• 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 does not take this approach because its consumer lies outside the kernel and it must support cross-node communication, so blocks arriving early bring it no benefit. Both designs are sensible under their own constraints.

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

• DeepGEMM: PR #304 (introduces Mega MoE), PR #316 (official benchmarks)
• Cursor, Mixture-of-Kittens: our open-source MoE megakernel for NVL72s (2026-08)
• 木子知, DeepGEMM Mega MOE算子详解 (in Chinese)