How GPUs Fuse Communication and Computation During LLM Inference

If you have read about how a large Mixture-of-Experts (MoE) model is served on several GPUs, you have probably seen the MoE layer summarized in three words: dispatch, experts, combine. We send the tokens to the GPUs that hold their experts, compute the experts there, and send the results back. On the GPU, these are separate programs (kernels) that run one after the other.

Let’s say we serve DeepSeek-V4-Flash on a machine with 8 B300 GPUs of NVIDIA’s Blackwell generation. Each GPU brings 512 tokens to a decode step, for example 128 requests with 4 tokens each from speculative decoding. The model has 43 layers, and each of them runs this chain of kernels. On our machine, the chain takes {{chain_us_512}} microseconds per layer. A newer kernel called Mega MoE produces the same output, bit for bit, in {{fused_us_512}} microseconds. Over the 43 layers, that would save {{step_saved_ms}} milliseconds per decode step if a whole server behaved like this benchmark.

Mega MoE is part of DeepGEMM, the kernel library that DeepSeek open-sourced, and its README describes it in one sentence:

Mega MoE fuses and overlaps EP dispatch, linear 1 and linear 2 (FP8xFP4 or FP8xFP8), SwiGLU, and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation.

EP stands for expert parallelism here, and linear 1 and 2 are the two matrix multiplications of an expert. There is one large kernel per GPU, and the 8 of them exchange tokens over NVLink, the direct links between the GPUs of a machine, while they compute. This is what the DeepSeek-V4 recipes in the SGLang cookbook and in the open InferenceX benchmark turn on for high throughput. In return, it needs a buffer of several gigabytes on every GPU and restricts how we can run the server.

So, if the work is the same, where does the saved time come from? In this article, I want to follow one layer through the kernel, ask at every step what it has to wait for, and then look at what it costs.

My descriptions come from reading the source code of DeepGEMM, SGLang (the inference engine), and FlashInfer (the kernel library it builds on), and most concrete numbers are for DeepSeek-V4-Flash. The measurements are my own, from one machine with 8 B300 GPUs. This article continues a series that started with How GPUs Talk to Each Other During LLM Inference, but you don’t need to have read the earlier ones.

1. An MoE Layer as a Chain of Kernels#

1.1 One Token Through the MoE Module#

An MoE model replaces the feed-forward module of a transformer layer with many small ones, called experts, and a router picks a few of them for each token. DeepSeek-V4-Flash has 256 experts per layer, and the router picks 6 for each token, together with a weight for each.

An expert consists of two matrix multiplications with an activation in between. The first one projects the 4,096 hidden values of a token to 2 × 2,048 values, called gate and up. The activation, SwiGLU, turns these into 2,048 values, silu(gate) * up, and the second multiplication projects them back to 4,096. The output of the module is the sum of the 6 expert outputs, each multiplied by its weight.

A kernel that multiplies matrices is called a GEMM, and it runs on the GPU’s tensor cores, the units built for this. A GEMM multiplies a set of rows by one matrix, but in an MoE layer every token goes to different experts. So, we first make one (token, expert) pair for each expert a token was routed to and sort the pairs by expert. A grouped GEMM then multiplies each expert’s rows by that expert’s matrix. At the end, the rows are put back into token order and summed, as shown in Figure 1 below.

The part of an inference engine that does what Figure 1 shows is called an MoE runner. Its stages are permute for the sort, GEMM 1, activation, GEMM 2, and finalize, which puts the rows back, weights them, and sums them.

One more detail: these matrices don’t hold the usual 16-bit floats. In DeepSeek-V4-Flash, an expert weight takes 4 bits (FP4) and a token value that enters a GEMM takes 8 bits (FP8), and each group of 32 values shares a scale that restores their real size. The tensor cores of a B300 multiply such numbers directly, but what comes out is an ordinary float again. So, the output of GEMM 1 has to be turned back into 8-bit values with new scales before GEMM 2 can read it, and the activation step in the middle is really “SwiGLU and re-quantize.”

1.2 Expert Parallelism Adds Two Steps#

Everything so far happened on one GPU. With 8 GPUs, the setups that aim for high throughput use expert parallelism for the MoE module, which means that each GPU holds a subset of the experts, each one complete. With 256 experts, that’s 32 experts per GPU.

In these setups, every GPU runs the attention module for its own share of the requests, which SGLang calls DP attention. When a layer reaches its MoE module, each GPU holds only the tokens of its own requests, the 512 of our example.

Since a token’s 6 experts usually sit on several GPUs, the token has to travel to them, and the 6 results have to travel back. This adds one step at each end of the runner:

• Dispatch: every (token, expert) pair is sent to the GPU that holds the expert.
• Combine: every expert output is sent back to the GPU of its token and summed there.

In our example, 512 tokens × 8 GPUs × 6 experts gives 24,576 pairs per layer, or 96 per expert on average. In Mega MoE, each pair travels as 4,096 bytes of FP8 values, 128 bytes of scales, and a 4-byte weight, 4,228 bytes in total. Its 16-bit result travels back as 8,192 bytes.

As the baseline for this article, I use the chain that DeepGEMM’s own test compares Mega MoE against. Dispatch and combine come from version 2 of DeepEP, DeepSeek’s communication library. There, the sending GPU writes each token into a buffer on the receiving GPU, and a small kernel on the receiving side then copies the tokens into expert order. The two grouped GEMMs come from DeepGEMM, and one more kernel does the activation and the re-quantization. Figure 2 below shows this chain on 2 GPUs.

Note the two strips at the bottom of Figure 2: the tensor cores are idle while tokens and results travel, and NVLink is idle during both GEMMs. During the activation, neither has anything to do.

1.3 Where the Time Goes#

So, how much of the chain is communication? I ran DeepGEMM’s test with the sizes of DeepSeek-V4-Flash on the 8 B300 GPUs and recorded the GPU time of every kernel in the chain. The routing in this test is random, which means that every expert gets about the same number of tokens. Figure 3 below shows the result for 1, 512, and 8,192 tokens per GPU.

[Figure 3 is not drawn yet]

{{fig3_reading}}

1.4 Summary#

To summarize, an expert-parallel MoE layer has five steps: dispatch, GEMM 1, activation, GEMM 2, and combine. In the chain, each of them is at least one kernel of its own. (The test’s chain launches {{chain_launches}} kernels per layer, and DeepGEMM’s release notes speak of 11 kernels that became 1.)

For the rest of this article, I will ask the same question at every step: what does it wait for before it can start? For the chain, the answer is always “the whole step before it.” Mega MoE replaces the runner and both communication steps at once, and this changes the answer for almost every step.

2. Two Things the Fused Kernel Is Built From#

What would a single kernel need to replace the whole chain? It would have to do several things at the same time and reach the memory of the other GPUs.

2.1 One Thread Block, Several Jobs#

Every launch of a kernel costs some microseconds before any useful work starts. A kernel’s threads are grouped into thread blocks, and each of them runs on one streaming multiprocessor (SM). A B300 has 148 SMs. Inside a thread block, every 32 threads form a warp.

A modern GEMM kernel launches one thread block per SM, and each stays alive and works through the matrices one rectangular piece, or tile, at a time. Its warps have different jobs:

• One loader warp per matrix asks the SM’s bulk-copy unit to bring the next tile from the GPU’s main memory into the SM’s small and fast shared memory.
• A multiplier warp tells the tensor cores to multiply the two tiles, and the product accumulates in a scratch area of the SM called tensor memory.
• The epilogue warps read the finished product out of tensor memory, post-process it, and store it to main memory.

Each job hands its work to the next with a small signal that says “ready.” Since the loaders can already fetch the next tiles while the multiplier is busy with the current one, the three jobs overlap in time. Figure 4 below shows the path of the data.

The result box in Figure 4 is one rectangle of the output: some token rows times some output columns. Computing one such rectangle is what I will call a task. A task takes several steps, because each output value is a sum over the dimension that the two matrices share. In each step, the loaders fetch one tile of each matrix along that dimension, and the multiplier adds the product onto the same spot in tensor memory. The epilogue then runs once, after the last step. In the DeepGEMM kernels below, 2 SMs share a task as an SM pair and split its columns between them.

2.2 Memory That Every GPU Can Reach#

Inside a machine with NVLink, a kernel on one GPU can read and write memory on another GPU with ordinary load and store instructions, as long as that memory has been mapped beforehand. Usually, every GPU allocates a buffer with the same layout, and each GPU’s process maps the buffers of all the others, which is often called symmetric memory.

Since all buffers have the same layout, a pointer into my buffer becomes a pointer to the same place in your buffer once I add your GPU’s offset. In DeepGEMM, that’s a single line (simplified from here):

// Simplified from SymBuffer::map
ptr_t map(ptr_t ptr, uint32_t dst_gpu) {
    return ptr + offsets[dst_gpu];   // the same place, in the buffer of dst_gpu
}

A load or a store through such a pointer travels over NVLink, so it only works inside an NVLink domain, the set of GPUs that reach each other this way. In our case, that’s the 8 GPUs of the machine.

What this kind of memory doesn’t give us is a way to notify another GPU. If GPU 0 writes something into GPU 1’s buffer, nothing on GPU 1 wakes up. That’s why SM pairs and GPUs signal each other only with counters. A counter is one word in the buffer. One thread adds to it, and another reads it in a loop until it sees the value it is waiting for. This is called spinning, and every wait below is built this way.

Mega MoE uses both ideas: the thread block of Figure 4 gets one more job, fetching tokens from the other GPUs, and the layer’s tokens and results live in a symmetric buffer.

3. Mega MoE, Step by Step#

Let’s now follow the tokens of one MoE layer from the moment they enter the fused kernel until the summed result leaves it.

The kernel that I describe is the one that SGLang runs for DeepSeek-V4 on Blackwell GPUs (source). SGLang ships its own fork of DeepGEMM, which is based on DeepSeek’s release from September 10, 2026, so a few details may differ from the latest upstream version.

To have something to point at, I will use a small example in the figures: 2 GPUs with 4 tokens each, 2 experts per token, and 2 experts per GPU.

3.1 Before the Kernel: The Buffer and the Cast Kernel#

Everything that the layer touches, except the expert weights and the final output, lives in the symmetric buffer. The buffer of each GPU is divided into regions, and five of them matter for us:

• The input region holds this GPU’s own tokens in FP8, with their scales, their 6 expert numbers, and their 6 weights.
• The tables and counters hold the bookkeeping.
• The arrived region will hold the tokens that this GPU fetches for its own experts.
• The intermediate region will hold what goes from GEMM 1 to GEMM 2.
• The result slots have room for 6 expert outputs for each token of this GPU.

Filling the input region is the one thing that happens before the fused kernel. In SGLang, a small cast kernel converts the layer input to FP8, computes the scales, and writes both, together with the router’s output, into the input region. SGLang launches the fused kernel right behind it, so its launch and setup overlap with the end of the cast, and it then waits for the cast to finish. Figure 5 below shows the buffers at this moment in the small example.

The fused kernel runs one thread block on almost every SM, in SM pairs. Each thread block has 16 warps: 4 new dispatch warps, 8 epilogue warps, and 4 GEMM warps. Those 4 are the token loader, the weight loader, the multiplier, and one that claims tasks for the SM pair. The boxes at the bottom of each GPU in Figure 5 are these three groups of warps.

3.2 Dispatch Becomes a Pull#

In the chain, dispatch is a push: the GPU that has a token writes it into a buffer on the GPU that needs it. Mega MoE turns this around, so that the GPU that has the token only announces it and the GPU that needs it fetches it.

The first part is the announcement. The dispatch warps of each GPU count how many (token, expert) pairs go to each of the 256 experts. Then, for every pair, they write a 4-byte index into a table on the GPU that holds the expert. The index says which token it is and which of its 6 choices, and it goes into the part of the table that is reserved for this expert and this source GPU.

The second part is a barrier, a point that nobody passes before everybody has reached it. We need it because a GPU can only place the rows of its experts once it knows how many rows each expert gets from each source, and the counts are only complete once every GPU has announced. It has two levels. Inside a GPU, every thread block adds to a counter and spins until all thread blocks of that GPU have added, so the GPU’s own counts are final. One thread block then writes these counts into the buffers of the GPUs that hold the experts. Then, across GPUs, that thread block adds to a second counter on all 8 GPUs and spins until the one on its own GPU shows all 8 have added, while the other thread blocks wait for it at one more barrier inside the GPU. By then, all counts have arrived and every table is complete.

The third part is the fetch. Here, the dispatch warps of each GPU go through the rows of their own experts, as shown below (simplified):

// Simplified from the dispatch warps (tma = the bulk-copy unit, smem = shared memory)
for (row in the rows of my experts that this warp handles) {
    src_gpu = the source GPU of this row, round-robin;  // computed from the row number
    idx     = index_table[expert][src_gpu][next];           // written by src_gpu
    token   = idx / 6;  choice = idx % 6;                   // idx 13 is choice 1 of token 2

    tma_load(smem, map(input[token], src_gpu));       // 4,096 bytes over NVLink
    tma_store(arrived[row], smem);
    origin[row] = {src_gpu, token, choice};                 // for the way back

    add(block_counter[row / BLOCK_M], 1);                   // the expert's last row adds the rest
}

A GPU takes the sources of an expert in turn, which is why the rows alternate between blue and orange in the figure below. The scales and the routing weight follow with plain loads. Figure 6 below shows the result in the small example.

In Figure 6, the rows land in the arrived region already sorted by expert, because the fetching GPU decides where each row goes. So, the permute stage has moved into the copy, at the price of one copy per pair. If two experts of a token sit on the same GPU, DeepEP sends the token there once. Mega MoE fetches it twice, like a3 on GPU 1 in Figure 6, since each expert needs the row among its own rows. With uniform routing on 8 GPUs, the 6 experts of a token sit on 8 × (1 − (7/8)⁶) ≈ 4.4 different GPUs on average, which means that Mega MoE moves 6 copies where 4.4 would be enough.

Why does Mega MoE pull? For once, we don’t have to guess: DeepSeek’s technical report on V4 says it in a list of lessons for hardware designers:

In the dispatch stage, we adopt a pull-based approach where each GPU actively reads activations from remote GPUs, avoiding the high notification latency that fine-grained push entails.

This is the missing notification again. If GPU 0 pushed single tokens to GPU 1, GPU 1 would have to keep checking which of them have arrived. When GPU 1 fetches them itself, it knows that a row is complete as soon as its own store has finished.

So, dispatch waits for the cast kernel, the counts of all GPUs, and one barrier, and nothing after that.

3.3 GEMM 1 Starts as Soon as an Expert’s Tokens Are There#

The last line of the fetch loop is where dispatch hands over to the GEMMs. The arrived region is divided into row blocks, and each belongs to one expert and has a counter. The kernel picks one block size per launch. Each fetched row adds one to the counter of its row block, and the last row of an expert fills the counter up to the block size, so a full counter always means that nothing more will come.

On the GEMM side, no SM pair has a fixed share of the experts. Instead, the work is cut into tasks, and a GEMM 1 task is one row block times 256 of the 4,096 gate and up columns, which makes 16 tasks per row block. Each task goes through the 4,096 input values of its rows in 32 steps of 128.

In each SM pair, one warp claims the next task by adding one to a counter that all SM pairs share, and passes the task to the pair’s other warps. The token loader of a GEMM 1 task then spins on one thing only, the counter of its row block:

// Simplified: the token loader of a GEMM 1 task
while (load(block_counter[task.block]) != BLOCK_M) { /* spin */ }

So, a claimed task starts as soon as its own row block is there, however many other rows are still on their way. Tasks are claimed in expert order, which is also the order of the fetch. Figure 7 below follows one row block.

In the upper half of Figure 7, the counter says full, so all 16 GEMM 1 tasks may start: 6 are stored, 3 are being multiplied, and 7 are claimed and queued.

How big is a row block? In our decode example, an expert’s 96 rows fit into a single row block of 128, so fetching and multiplying overlap only across experts. With 8,192 tokens per GPU, as in a prefill chunk, an expert gets about 1,536 rows in row blocks of 240, which become ready one after the other.

GEMM 1, then, waits for one row block of one expert.

3.4 The Cost of a Row Block#

If GEMM 1 can start as soon as a row block is full, wouldn’t smaller row blocks let it start earlier? They would, but every row block has a cost, and that cost shows where the time of this kernel goes.

To multiply a row block, the loaders have to read all the weights of the expert, no matter how few rows it has. The 4-bit weights of one expert take 12 MB, so in our decode example, a GPU reads 32 × 12 MB = 384 MB of weights per layer. (I am using MB and GB for powers of 1,024 throughout.) Compare this with the tokens: the 3,072 incoming rows are 12 MB, and the 3,072 outgoing results are 24 MB. That is about 11 times as many distinct bytes of weights as of tokens.

To check this, I ran the fused kernel on a single GPU with 32 experts, so that it does the work of one of our 8 GPUs while all the fetches stay local. It takes 40 microseconds with 1 token (6 touched experts), 68 with 4 tokens (about 18), 95 with 16 tokens (about 31), and 99 with 64 tokens (all 32). From there to 512 tokens, it only grows to 112. So, the time depends on the touched experts and hardly on the rows: each further expert costs about 2.3 microseconds, on top of about 27 that don’t depend on the batch. This doesn’t prove that reading the weights is the cost, since every touched expert also adds tasks with fixed work, and I haven’t profiled the memory traffic. But at the 6.5 TB/s that DeepGEMM’s test assumes for memory, reading the 12 MB of one expert would take 1.9 microseconds.

The code that picks the block size says so in a comment, where an M block is a row block: “Every M block costs roughly 128 extra rows (its tasks reload the weight tiles), so use the fewest blocks per expert and the smallest tile covering them […].”

3.5 SwiGLU, the Routing Weight, and Re-Quantization in One Epilogue#

In the chain, the activation is a kernel of its own that reads the output of GEMM 1 from memory. In Mega MoE, it is the epilogue of the GEMM 1 task: the epilogue warps take the product out of tensor memory and do the whole activation step before they store anything.

For this to work, a gate value and its up value have to come out of the same task, but the checkpoint stores the 2,048 gate columns of the first matrix before the 2,048 up columns. So, SGLang reorders them when it loads the model: 8 gate columns, then the matching 8 up columns, then the next 8 gate columns, and so on. A task of 256 columns then holds 128 gate values with their up values, and the epilogue becomes short (simplified):

// Simplified from the GEMM 1 epilogue
gate  = min(gate, limit);  up = clamp(up, -limit, limit);
act   = gate * sigmoid(gate) * up;         // SwiGLU
act   = act * routing_weight[row];         // the router's weight for this pair
scale = pow2_at_least(max_abs(act over a group of 32) / 448);   // 448 = largest FP8 value
store(intermediate, to_fp8(act / scale), scale);

The first line is a clamp that DeepSeek-V4 applies before SwiGLU, and the last two lines are the re-quantization. The third line applies the routing weight, which mathematically belongs to the final sum. But GEMM 2 is linear, so multiplying its input by the weight gives the same result as multiplying its output. Doing it here means that the results can later be summed without any weights. The activation kernel of the chain applies the weight at the same place, which is why the two outputs can agree bit for bit.

SGLang’s Triton runner has a similar reordering as an opt-in for unquantized weights. It isn’t the default there, because weight updates and moving experts between GPUs at run time (EPLB) expect the layout of the checkpoint.

The activation, then, waits for nothing of its own, since it is the tail of its GEMM 1 task. The multiplier, though, has to wait for the epilogues whenever they fall behind, since tensor memory holds two products.

3.6 GEMM 2 Starts Before GEMM 1 Is Finished#

GEMM 2 multiplies the 2,048 activation values of each row by the second weight matrix, and a GEMM 2 task is again one row block times 256 of the 4,096 output columns. GEMM 1’s output columns are now the dimension that GEMM 2 sums over, so every GEMM 2 task goes through the 2,048 values in 16 steps of 128. These values come from the 16 GEMM 1 tasks of the same row block, since the 256 gate and up columns of one GEMM 1 task turn into 128 activation values. So, step k of a GEMM 2 task reads what GEMM 1 task k has written.

The kernel uses this by keeping one mask per row block, with two bits per GEMM 1 task, one for each SM that ran it. An SM sets its bit when its half of the task is stored. Before a GEMM 2 task takes step k, its loader checks the two bits of that step. If they are set, it goes ahead, even if other GEMM 1 tasks of the row block are still running or queued in their SM pair. This is the lower half of Figure 7 above.

If all SM pairs claimed GEMM 2 tasks, nobody would be left to run the GEMM 1 tasks they are waiting for. To rule this out, the warp that claims tasks starts with enough rounds of GEMM 1 tasks and then alternates between the two kinds.

GEMM 2, then, waits for one GEMM 1 task per step of its input, and all of this runs while the dispatch warps of the same thread blocks are still fetching tokens.

3.7 Combine: Write to the Source, Sum at the End#

At the end of a GEMM 2 task, the epilogue converts each of its rows to 16 bits and stores it directly into a result slot on the GPU that the token came from. The address comes from the note that the dispatch warps made during the fetch: the GPU, the token, and which of the token’s 6 choices this expert was. This way, the output for choice j of a token always lands in slot j of that token. This is a push, which dispatch avoided, but here nobody waits for a single row, because only the sum at the end reads the slots.

With this, the finalize stage of the runner is gone as well, since the 6 slots of a token only have to be added up. (SGLang’s TensorRT-LLM runner shortens this stage in its own way: at small batches, one kernel un-sorts and sums.)

This addition is the one step that still comes at the end. When a GPU has finished all of its GEMM 2 tasks, it marks itself as done on every GPU. The epilogue warps of each GPU then go through their own tokens and wait, for each one, until the GPUs that hold its 6 experts are marked as done. After that, they load the 6 slots one after the other, add them up as 32-bit floats in slot order, round once to 16 bits, and write the row to the output of the layer, as shown in Figure 8 below.

In Figure 8, a0 gets a result from each GPU, so its sum has to wait for both. What does combine wait for? The stores wait for nothing, since they are the tail of a GEMM 2 task. The sum, though, waits for whole GPUs, its own and those of the token’s experts, which makes it a coarse wait.

During the sum, the dispatch warps reset the counters and masks for the next layer, and a second barrier across the GPUs ends the kernel.

3.8 The Shared Expert#

One stretch is left in which the tensor cores have nothing to do: the beginning of the kernel, with the counting, the announcing, the barrier, and the first fetches. Mega MoE uses it for a part of the MoE module that I have left aside so far, the shared expert.

DeepSeek-V4-Flash has one shared expert per layer. Every token goes through it without any routing, and its output is added to the sum of the 6 routed ones. Outside the fused kernel, SGLang runs it as a separate dense module, on a different stream than the routed experts. A stream is a queue of kernels on the same GPU. SGLang also has what it calls shared-expert fusion. There, the shared expert becomes one more expert that every token selects in addition to its 6, and every GPU keeps a copy so that this pair never travels. For DeepSeek-V4, fusion is off by default, since the shared expert is stored in FP8 and the routed experts in FP4. An opt-in flag re-quantizes it to FP4 at load.

Inside the kernel, the shared expert is two more kinds of tasks. The first reads its rows directly from the input region and needs no counts, so the task-claiming warp takes it first. Its multiplication waits for nothing but the cast kernel, and its epilogue runs after the first barrier. The second writes into a 7th result slot.

Running it inside the kernel is SGLang’s default, although the pull request that added this measured only about 1% more decode throughput on 4 B300 GPUs. It also explains why: on separate streams, the shared expert already overlapped almost completely with the routed kernel.

3.9 The Whole Layer Again#

We can now redraw the timeline from Figure 2 for the fused kernel, as shown in Figure 9 below.

From left to right, Figure 9 starts with a short stretch where every GPU counts and announces its pairs while the tensor cores start on the shared expert. Then comes the first barrier. After it, the fetches, both GEMMs, the activation, and the stores on the other GPUs all run at the same time, and the sum, the reset, and the second barrier come at the end. In this drawing, the two strips at the bottom overlap almost everywhere.

To summarize, this is what each step waits for inside Mega MoE:

• Dispatch waits for the cast kernel, the counts of all GPUs, and one barrier, and nothing after that.
• GEMM 1 waits for one row block of one expert.
• The activation waits for nothing, since it is the tail of its GEMM 1 task.
• GEMM 2 waits for one GEMM 1 task per step of its input.
• In combine, the stores wait for nothing, and the sum waits for its own GPU and the GPUs that hold a token’s experts.

How much of the chain’s time does this remove?

4. How Much Faster Is the Layer?#

I ran the fused kernel and the chain side by side, for 1 to 8,192 tokens per GPU.

Both runs use the sizes of DeepSeek-V4-Flash with random weights and random routing on the 8 B300 GPUs. I left the shared expert out and gave the kernel all 148 SMs here. Before timing, I compared the two outputs, which are bit-identical at every batch size. The times are the GPU times that the PyTorch profiler reports for the kernels, averaged over the 8 GPUs, which a barrier lines up before each run. They don’t include the copies that put the inputs into the input region. {{method_extra}}

Figure 10 below shows the result.

[Figure 10 is not drawn yet]

{{fig10_reading}}

For comparison, DeepSeek reports 1.96× at 1 token, 1.73× at 512, and 1.56× at 8,192 tokens per GPU for the same sizes, in the pull request with their benchmarks. Their numbers are from April 2026, which is four upstream updates before the version I ran, and they don’t name the hardware.

{{sec4_variants}}

5. What It Costs#

The first cost is the buffer, where every step in section 3 found a region that was already there, because a kernel can’t ask the CPU for memory in the middle of a layer.

5.1 Sizing the Buffer#

How large does the buffer have to be? Since nothing can be added to it later, the kernel is compiled for a maximum number of tokens per GPU, which I will call the cap. SGLang’s default is 8,192, and DeepGEMM rounds it up to 9,600, the next multiple of 1,920.

The input region and the result slots are easy to size, because they belong to this GPU’s own tokens. The arrived region is the hard one, since how many rows arrive at a GPU depends on the router. In the worst case, every token of every GPU sends all 6 of its copies to the same GPU, which would be 8 × 9,600 × 6 = 460,800 rows. SGLang’s fork sizes the arrived region and the intermediate region for that case, plus a little padding.

For DeepSeek-V4-Flash on 8 GPUs, this adds up to 4.5 GB per GPU for the routed experts, and to 4.7 GB with the shared expert inside the kernel. All 43 layers share this one buffer. This is from the layout code, and {{buffer_check}}. SGLang allocates it on the first forward pass, outside its own memory accounting, which is why the cookbook’s recipes lower --mem-fraction-static as the cap grows. The largest part is the arrived tokens: 460,800 rows × 4,096 bytes is 1.8 GB, and the intermediate values, at 2,048 bytes per row, are 0.9 GB. The scales would be 1/32 of both, but each row block’s scales are padded to 128 rows. Since the layout must allow row blocks as small as 8 rows, the padding costs a factor of 16 and brings them to 1.3 GB. The result slots, 9,600 × 6 × 8,192 bytes, take 0.4 GB.

Upstream DeepGEMM doesn’t size for the worst case anymore. It keeps the arrived and intermediate regions as a ring of a few thousand rows, enough for the row blocks that can be in use at once. It reuses a row as soon as the next step has read it, which adds two waits to steps that used to wait for nothing. A dispatch warp spins until GEMM 1 has read what was in its row before, and a GEMM 1 epilogue does the same for its row of the intermediate region.

SGLang’s fork uses the ring only if the rounded cap is 5,760 or less. According to a comment in its code, larger caps keep the full-size regions to preserve a tuned prefill path that never wraps around. With a cap of 4,096, which rounds up to 5,760, the whole buffer shrinks to 0.35 GB, as shown in Figure 11 below.

As we can see in Figure 11, the small buffer consists mostly of result slots, which still need room for 6 outputs for every token up to the rounded cap. {{ring_timing}}

5.2 The Token Budget#

Since the buffer is sized before the first request, the cap is a hard limit for the kernel. Above it, SGLang quietly takes the unfused MoE path, which models such as DeepSeek-V4 don’t have. The engine has to rule this out in advance, so at startup SGLang tries to bound how many tokens one GPU can see in one forward pass:

decode:   largest batch in the CUDA graph  ×  tokens per request
prefill:  chunk size  ÷  number of DP attention groups

For decoding, SGLang records the kernels of a decode step once per batch size as a CUDA graph and replays them, so the check takes the largest recorded batch. With speculative decoding, it multiplies by the draft tokens that each request brings to a step. For prefill, SGLang feeds a long prompt through in chunks of at most --chunked-prefill-size tokens, and with DP attention the whole server shares this budget, so it is divided among the GPUs. For the checked models, SGLang refuses to start with --chunked-prefill-size -1, which turns chunking off, since a prompt of any length would then go through in one piece and no cap could be guaranteed.

For example, the InferenceX recipe for DeepSeek-V4 on 8 B300 GPUs uses a chunk size of 65,536 with DP attention on all 8 GPUs, which gives 8,192 tokens per GPU. It sets the cap to 8,320, which is rounded up to 9,600 like the default. The cookbook’s balanced recipe for DeepSeek-V4-Pro on B200 and SGLang’s own test for DeepSeek-V4-Flash use a cap of 4,096 instead, small enough for the ring.

The check runs for nine architectures, DeepSeek-V4 and the Qwen MoE family among them, and for every checkpoint with expert weights in NVFP4, NVIDIA’s 4-bit format, because with megamoe they have no other MoE path.

5.3 The Barriers Need the Whole GPU#

The next cost comes from the barriers, which open only when every thread block has added to a counter, so all of them have to run at once. Below is the core of the one inside a GPU (simplified from here):

// Simplified from grid_sync: a barrier across all thread blocks of one GPU
old = atomic_add(counter, is_first_thread_block ? FLAG - (num_thread_blocks - 1) : 1);
do {
    now = load(counter);
} while (((now ^ old) & FLAG) == 0);     // spin until the top bit has flipped

The first thread block adds a large flag value minus the number of the others, so the flag bit flips only once all of them have added.

What if a kernel on another stream occupies an SM at that moment? Then one SM pair of Mega MoE can’t start, and the thread blocks that did start spin on a counter that can’t reach its value. Since every other GPU waits for this GPU’s counts, the layer stalls on all of them for as long as the other kernel runs.

SGLang ran into this as a timeout on GB200 machines, and the fix was to launch Mega MoE on 146 of the 148 SMs, which leaves 2 SMs for a small kernel on another stream. Anything larger still stalls the layer.

The barrier across GPUs needs every GPU as well, so each has to launch the kernel for every layer, even with no tokens, because the others wait for its counts.

The two barriers and the wait before the sum give up after about 300 seconds with an assertion that kills the kernel. The inner waits never give up: for the counts, for a row block’s counter, for the mask, and for a row of the ring.

5.4 Shrinking the Result Slots#

The result slots exist because combine pushes every result to the token’s GPU and sums only at the end. Most of the traffic goes into them, since the 16-bit result is 8,192 of the 4,228 + 8,192 = 12,420 bytes that cross NVLink for each pair. That’s why the newer variants of the kernel try to shrink them. The most interesting variants are in FlashInfer, which contains a second implementation of Mega MoE written by NVIDIA.

The first idea is to add on arrival. Since the routing weight is already inside the result, a row doesn’t have to be stored in a slot at all. The GEMM 2 epilogue can add it straight into the token’s output row on the other GPU with an atomic add across NVLink, and when all GPUs are done, the output is complete. The slots, the separate sum, and its wait for whole GPUs all disappear.

This changes the result a little, though. With slots, the 6 values are added in a fixed order as 32-bit floats and rounded once. With atomic adds, they are added in whatever order the rows arrive from up to 6 GPUs, and every addition rounds to 16 bits. Floating-point addition can give different results in different orders, as Figure 12 below shows.

In Figure 12, adding 256 + 1 + 1 from left to right rounds back to 256 at each step, while 1 + 1 + 256 gives 258. As a result, the output of a layer can differ in its last bits from one run to the next. FlashInfer’s configuration file says this about the option: “~1-2% faster and removes the multi-GB combine staging from the symmetric workspace […] accumulation order is nondeterministic”. SGLang has it as an environment variable that is off by default. It only affects FlashInfer’s own kernels, so it doesn’t reach DeepSeek-V4-Flash, whose FP4 experts run in the DeepGEMM kernel.

The second idea is to send fewer bits. FlashInfer’s NVFP4 kernel can quantize the results to 8 or 4 bits with scales before they travel, which halves or quarters the result bytes, and the receiving GPU converts them back before it sums. According to the same file, this trades “accuracy for NVLink bandwidth”. It can’t be combined with adding on arrival, I think because an atomic add happens in memory, so no code on the receiving GPU could convert the values back first. SGLang’s DeepGEMM fork has an 8-bit variant as well, which we turn on with the library’s variable DG_USE_FP8_COMBINE, and that one applies to our example. {{fp8_combine_result}}

5.5 When Not to Use It#

These costs fit a high-throughput server with DP attention, where batches are large and every GPU has tokens in every layer. Mega MoE also needs expert parallelism across all GPUs of the server. Since the kernel reads the other GPUs’ buffers with plain loads, I would also expect them all to share one NVLink domain, which in a GB300 rack spans several machines. The cookbook lists a Kimi K3 recipe with megamoe on four 8-GPU machines, though, which I haven’t tried and can’t explain. SGLang also rejects its two-batch and single-batch overlap features, because they need separate dispatch and combine steps, which Mega MoE doesn’t have. In the SGLang cookbook’s recipes for DeepSeek-V4-Flash, the balanced one uses DeepEP, and only the high-throughput one uses Mega MoE.

Even there, it isn’t always faster. The pull request that brought Mega MoE to the Qwen MoE models compared it with FlashInfer’s separate kernels on a 397B model on GB300 machines. It was up to 7% behind on a single machine and mostly 1% to 4% ahead across 2 to 4 machines, and the author called it “backend coverage rather than a speedup.”

{{skew_note}}

6. Mega MoE in SGLang#

6.1 Communication Backend and Runner#

SGLang builds an MoE layer from two interchangeable parts, and Mega MoE is both of them at once. The first part is the communication backend, which provides dispatch and combine, and the second is the runner backend, which provides the stages in between. We select them with --moe-a2a-backend (a2a stands for all-to-all) and --moe-runner-backend, and the layer’s forward pass then consists of three calls (simplified from here):

# Simplified from FusedMoE.forward_impl
dispatched = self.dispatcher.dispatch(hidden_states, topk_output)   # communication backend
computed   = self.run_moe_core(dispatched)                          # runner backend
return self.dispatcher.combine(computed)                            # communication backend

There are more than a dozen runner backends, because the right grouped GEMM depends on the format of the expert weights and on the GPU generation. SGLang also picks one this way if we don’t name it. For example, an FP8 checkpoint gets deep_gemm behind DeepEP and triton otherwise, and an NVFP4 checkpoint gets flashinfer_trtllm on data-center Blackwell GPUs and marlin on older ones.

One kernel can’t be spread over these three calls, so SGLang wires in its two implementations of Mega MoE differently, as shown in Figure 13 below.

In both of the lower rows of Figure 13, one kernel does all the work, and only the place where SGLang calls it differs.

6.2 The Two Implementations#

The first implementation is the DeepGEMM kernel from section 3. We select it as a communication backend, with --moe-a2a-backend megamoe, and the MoE module of the model then calls the kernel directly and skips the three calls above. The cookbook’s high-throughput command for the July checkpoint of DeepSeek-V4-Flash on 8 B200 GPUs is short:

SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=8320 \
sglang serve \
  --model-path deepseek-ai/DeepSeek-V4-Flash-0731 \
  --trust-remote-code \
  --tp 8 \
  --attn-dp-size 8 \
  --moe-a2a-backend megamoe

In this command, --tp 8 uses all 8 GPUs, the environment variable sets the cap (a rank is a GPU), and --attn-dp-size 8 turns on DP attention with every GPU as its own group. Choosing megamoe spreads the experts over all 8 GPUs automatically. The 2 reserved SMs and the shared expert inside the kernel are defaults with their own variables, SGLANG_OPT_DEEPGEMM_MEGA_MOE_RESERVED_SMS and SGLANG_OPT_DEEPGEMM_MEGA_MOE_FUSE_SHARED_EXPERTS. SGLang rewrites the alias --moe-runner-backend megamoe to the same setting.

The second implementation, flashinfer_megamoe, is FlashInfer’s kernel, written in CuTe DSL, a way of writing such kernels in Python. It works like the DeepGEMM kernel, with a pull for dispatch and stores for combine, and adds the combine variants. In SGLang, we select it with --moe-a2a-backend flashinfer_megamoe, and the runner of the same name is then chosen automatically.

Which of the two do we use, then? For FP4 experts as in DeepSeek-V4-Flash, flashinfer_megamoe wraps the DeepGEMM kernel, so both names end up in the same kernel, although flashinfer_megamoe has its own cap variable, reserves no SMs, and keeps the shared expert outside the kernel. Otherwise, the format of the expert weights mostly decides: megamoe also takes NVFP4 experts on Blackwell GPUs, and flashinfer_megamoe has its own kernels for NVFP4 and MXFP8, an 8-bit format with the same kind of scales. In any case, flashinfer_megamoe needs Blackwell GPUs and DP attention on every GPU, and it only accepts a fixed list of 10 architectures. It refuses checkpoints with plain FP8 expert weights, and the code doesn’t say why. My guess is that nobody has written the step that converts such weights into the layout of the kernel.

No cookbook recipe uses flashinfer_megamoe yet, and the kernel benchmarks that I found are mixed. In a FlashInfer pull request from July, on 4 GB200 GPUs with DeepSeek-V3 sizes, the NVFP4 kernel was on par with the DeepGEMM kernel of that time below 1,024 tokens per GPU and 1.6 times faster at 8,192. In another one from August, on 8 B200 GPUs with DeepSeek-V4-Flash sizes, the MXFP8 kernel took 261 microseconds against DeepGEMM’s 155 at 512 tokens per GPU, and the two were about equal at 8,192. The MXFP8 kernel reads twice the weight bytes, though, 8-bit against 4-bit.

6.3 A Whole Server#

{{server_ab}}

7. What I Left Out#

I skipped how DeepEP and the other communication backends work inside and how other engines such as vLLM wire in the same kernels. I also left out the balancing of load between experts and GPUs (EPLB and SGLang’s waterfill option, which works with megamoe) and the Hopper, AMD, and FP4-activation versions of the kernel.

8. Summary#

To wrap up, Figure 14 below collects what each step of the layer waits for, in the chain and in Mega MoE, together with who does it and where its data goes.

As Figure 14 shows, each thread block of the fused kernel has more jobs than a GEMM kernel’s, and the jobs wait for much smaller things. {{summary_speed}} We pay for this with a buffer that is sized before the first request, with limits on batch and chunk sizes, and with a kernel that needs the GPU to itself, since its barriers need every thread block.

Does this scale past 8 GPUs? The kernel accepts up to 72, the size of a GB300 rack. The worst case for the arrived region grows with the number of GPUs. If we split the 256 experts of DeepSeek-V4-Flash over 64 GPUs and keep the cap of 8,192 tokens per GPU, the buffer of SGLang’s fork comes to about 22 GB per GPU. If the fork used the ring at every cap, like upstream DeepGEMM, its layout code would give about 0.6 GB, which I haven’t measured.

DeepSeek’s report has two more lessons for hardware designers. The first is about the missing notification: “Future hardware with lower-latency cross-GPU signaling would make push viable and enable more natural communication patterns.” The second is about power: “Extreme kernel fusion drives compute, memory, and network to high load simultaneously, making power throttling a key performance limiter.” {{power_note}}