Attention, Parallelism, and Collective Communication in MoE Serving

Note: Significant portions of this article, including diagrams and code, were generated with the assistance of Claude (Anthropic). I wanted to be upfront about that, since I’d want to know before spending my time on something, and I imagine you might feel the same. If you find any inaccuracies feel free to point them out to me.

1. Introduction

This article covers the machinery that modern mixture-of-experts (MoE) serving systems are built on: the attention block, the parallelism strategies used to distribute it across devices, the routing and communication that MoE layers require, and the collective communication primitives underneath all of it. It was written while reading two recent systems papers, NanoCP and UltraEP, both of which assume this material as background.

The organizing fact is the following. Almost every operation in a transformer acts on each token independently. Two do not:

1. The attention score computation, which mixes information across token positions.
2. MoE dispatch, which sends tokens to experts that may reside on other devices.

Every communication cost discussed below arises from one of these two operations. The axis that a parallelism strategy partitions determines whether that strategy collides with one of them, and therefore whether it requires communication at all. Section 2 identifies the exception in attention, Section 3 works out its consequences for parallelism, Section 4 identifies the exception in MoE, and Section 5 gives the cost model that prices both.

Table 1: Notation used throughout.

2. The transformer layer

The input to an attention block is a tensor x of shape (B, T, d_model). Each of the B × T token positions holds a vector of length d_model, which is that token’s current representation at this depth in the network.

2.1 Projections

Three learned matrices, W_Q, W_K, and W_V, each of shape (d_model, d_model), produce the query, key, and value tensors:

Q = x @ W_Q,    K = x @ W_K,    V = x @ W_V

Each result has the same shape as the input, (B, T, d_model). None of these matrices is head-aware. The head structure does not exist yet at this point in the computation.

This step is per-token. Writing the matrix product elementwise,

(x @ W)[i, j] = Σ_c x[i, c] · W[c, j]

the row index i appears on both sides and is never summed over. Output row i depends on input row i and on W. Row i' of x, belonging to any other token, is never referenced. The independence holds because W is a fixed weight matrix, shared identically across all rows, so any operation of the form x @ W has the same property. This fact recurs throughout the article.

2.2 Head partitioning

Since d_model = n_heads × d_head by construction, the last axis of Q, K, and V can be reinterpreted as two axes, (n_heads, d_head). If d_model = 8 and n_heads = 2, a token’s eight-element key vector

[1.2, -0.4, 0.7, 2.1, 0.3, -1.5, 0.9, 1.8]

is reinterpreted as two four-element vectors, one per head:

head 0:  [1.2, -0.4, 0.7, 2.1]
head 1:  [0.3, -1.5, 0.9, 1.8]

No arithmetic occurs. The reshape is a reindexing of the same numbers, and in practice it does not move memory. The location of the boundary is a convention. What matters is that the same convention is applied to Q, K, and V, so that head 0’s query, key, and value slices correspond to one another.

2.3 The score computation

Within a single head, attention computes a matrix of scores between every query position and every key position:

S = Q Kᵀ / √d_head        shape (T, T)

A causal mask sets S[i, j] = -∞ for j > i, so that a query cannot attend to future positions. A softmax over the last axis converts each row into a distribution, and the output for each query is the corresponding weighted average of value vectors:

head_output = softmax(mask(S)) @ V        shape (T, d_head)

The head outputs are concatenated back to width d_model, multiplied by an output projection W_O, and added into the residual stream.

Figure 1: The attention block. Only the highlighted step mixes information across token positions.

2.4 The locus of cross-token interaction

The score computation is the only step in the block in which two distinct token positions interact. The difference lies in the operands:

x @ W        one operand (W) is shared by every row
Q @ Kᵀ       both operands are token-dependent

In the first product, the second operand is a fixed weight matrix, so the rows are independent. In the second, entry (i, j) of the result requires row i of Q and row j of K, and K is itself derived from every token’s hidden state. Computing the score row for query i therefore requires the key vectors of all positions j ≤ i.

Whether an operation’s operands are all token-independent determines every communication requirement in the rest of this article.

2.5 The feed-forward network

Each layer follows the attention block with a position-wise feed-forward network (FFN), which expands each token’s vector to a wider inner dimension, applies a nonlinearity, and contracts it back:

ffn_out = GELU(x @ W_up) @ W_down

with W_up of shape (d_model, d_ff) and W_down of shape (d_ff, d_model). The inner width d_ff is a capacity parameter, commonly a small multiple of d_model. The output width is fixed by the residual connection: ffn_out is added back into x, and addition requires matching shapes, so the FFN must return to width d_model.

Figure 2: The FFN expands, applies a nonlinearity, and contracts back to the residual width.

The nonlinearity is required for the wide inner layer to contribute anything. Matrix multiplication is associative, so without it

(x @ W_up) @ W_down  =  x @ (W_up @ W_down)

and the right side is a single (d_model, d_model) matrix. This is straightforward to confirm numerically: with GELU removed, the two-layer computation and the collapsed single-matrix computation agree to floating-point tolerance, and with GELU restored they do not.

Both x @ W_up and hidden @ W_down have the form x @ W, so by Section 2.4 the FFN is per-token and requires no communication across positions. What the FFN computes, as opposed to how it is shaped, is covered in Appendix B. Section 4 replaces this component with a mixture of experts.

3. Parallelism as partitioning of a tensor axis

After the head reshape, the activation tensor has four axes: (B, T, n_heads, d_head). Each parallelism strategy corresponds to partitioning one of these axes across devices. Communication is required exactly when the computation must combine information across slices of the partitioned axis.

Figure 3: The four axes of the attention activation tensor and the parallelism strategy associated with each.

Ordering the axes by the communication they require gives B < n_heads < T < d_head. The ordering does not follow the position of the axes in the tensor shape. It follows from whether attention must combine information across the axis.

3.1 Data parallelism partitions the batch

Attention is defined strictly within a sequence. No operation in the block combines information from two different batch elements. Partitioning B across devices therefore requires no communication: each device runs the complete attention computation, over all heads, for its own subset of sequences.

3.2 Tensor parallelism partitions the heads

Each head is computed in isolation from the others until the final output projection, so the head axis can be partitioned across devices. The mechanism operates on the weight matrices.

The columns of W_Q, W_K, and W_V are grouped by head: the first d_head output columns produce head 0, the next d_head produce head 1, and so on. Splitting these matrices column-wise assigns whole heads to devices. Each device computes its own heads’ queries, keys, and values, runs the entire score, softmax, and weighted-sum sequence locally, and never requires data from another device.

Figure 4: Column-wise partitioning of the projection weights assigns whole heads to devices. The KV cache partitions along with them.

The combine step requires an addition. Recombining per-head outputs looks like a concatenation, which across devices would be an all-gather. Megatron-style tensor parallelism instead folds the concatenation into the output projection: W_O is partitioned by rows, so each device multiplies its own head output by its own row block. Each device produces a tensor of full width d_model that represents a partial contribution to the correct answer. Splitting a matrix product along its contraction dimension makes the true result the sum of the partial products, so the combine is an addition, which across devices is an all-reduce.

Two consequences follow. TP requires the input x to be replicated across the group before the layer begins, which is typically already satisfied because the previous layer’s all-reduce broadcast it. And the same column-then-row pattern applies to the feed-forward network, with W_up partitioned by columns and W_down by rows, so a transformer layer under TP costs one all-reduce for attention and one for the FFN.

3.3 The exception: multi-head latent attention

The head-partitioning argument depends on each head owning a private slice of the KV cache. Multi-head latent attention (MLA), used by the DeepSeek model family, violates this assumption by construction.

Instead of caching per-head keys and values, MLA compresses them into a single shared latent vector per token, from which every head reconstructs its own keys and values via a per-head up-projection. This compression is what makes MLA’s KV cache small: DeepSeek-V3 caches a latent of width 512, plus a decoupled positional component of width 64, for 576 values per token, in place of full per-head keys and values.

Under head-wise TP, every device requires the entire latent to reconstruct even its own heads. There is no per-head slice to distribute, because the compression collapsed the head structure into one shared object. The cache therefore replicates instead of partitioning, and at a TP degree of 8 the model stores eight identical copies, which negates the compression that motivated MLA.

Figure 5: MLA’s shared latent has no head structure, so head-wise partitioning replicates the cache instead of dividing it.

MoE serving stacks built on MLA models, including both systems that motivated this article, therefore use data parallelism for attention, where each instance holds the small latent cache for its own requests with no sharing and no duplication. Expert parallelism, a different axis, is reserved for the feed-forward layers. This is the DP-EP configuration that recurs throughout the MoE serving literature.

3.4 Context parallelism partitions the sequence

Context parallelism divides a single sequence’s tokens across devices, where data parallelism divides whole sequences. The motivation is capacity: when one request’s KV cache approaches the memory of a single device, that cache must be distributed regardless of how the compute is arranged.

For the projection step, partitioning T is as free as partitioning B. The projections are per-token, so a device holding tokens 0 and 1 can compute their queries, keys, and values using the full, unpartitioned weight matrices with no communication, and likewise for a device holding tokens 2 and 3.

Figure 6: Under context parallelism the projections remain local. The score computation does not.

The score computation does not remain local. The query at position 3 must attend to keys at positions 0 through 3, half of which reside on the other device, so an exchange is unavoidable. Context parallelism is more expensive than the other two axes for this reason: its communication sits in the middle of the computation, where tensor parallelism’s sits only at the end.

Both implementations rely on the same property, which also underlies FlashAttention’s tiling and Flash-Decoding’s split reductions. A softmax-weighted average can be computed over disjoint blocks of keys and merged exactly, provided each partial result carries a running maximum and normalizer, equivalently a log-sum-exp term, alongside its unnormalized output. Given these, the partials recombine to the value the unpartitioned softmax would have produced.

The two implementations use the same merge and differ in topology:

• Query routing. The device holding the query sends it to whichever devices hold the relevant KV shards. Each computes a partial attention output and its log-sum-exp locally, and returns both. The originating device merges them. Helix and NanoCP both use this form.
• Ring rotation. KV chunks are rotated around a cycle of devices. At each hop, a device merges the arriving chunk into its running partial result and forwards its own chunk onward. After P hops every device has seen the entire sequence. This is Ring Attention.

Figure 7: Ring Attention rotates KV chunks around a fixed cycle instead of routing queries to fixed destinations.

Ring passes fixed-size chunks through a predictable, bandwidth-friendly pattern, which suits the large groups typical of training. Query routing sends data directly to its destination and avoids intermediate hops, which suits the tighter latency budgets and the smaller, more dynamic groups of decode-time serving.

3.5 The head dimension is not partitioned

Splitting d_head would divide a single dot product across devices, requiring an exchange of partial sums before the softmax could be applied to any score. The synchronization granularity is too fine to be practical, and the axis is left intact. Research on partitioning MLA’s latent dimension exists, motivated by the tensor-parallel limitation in Section 3.3, but it remains an open problem.

3.6 Pipeline parallelism partitions the layer stack

Pipeline parallelism does not partition the activation tensor. It assigns different layers to different devices, so one device holds layers 1 through 20 and passes activations onward to a device holding layers 21 through 40.

PP’s communication profile differs from TP’s. TP requires an all-reduce in every layer, which is frequent, latency-sensitive, synchronous traffic. PP requires one activation transfer per stage boundary. PP therefore tolerates slow interconnects far better, and it is the preferred choice once a deployment spans more devices than fit in a single high-bandwidth domain. Sarathi-Serve reports roughly a factor of two lower median latency for PP over TP when serving across nodes connected by commodity Ethernet.

PP pays for this with pipeline bubbles: idle stages waiting for work to arrive from upstream, which is most acute under low request load. The common production configuration follows from these two facts. TP is used within a fast interconnect domain, PP across domains, with DP and EP layered on top.

4. Mixture of experts

An MoE layer replaces the single shared FFN of Section 2.5 with E smaller ones, called experts, together with a gate that selects, per token, which k of them will process it.

4.1 Gating

The gate is a linear map from the token’s hidden state to E scores, followed by a top-k selection and a softmax restricted to the selected experts. Like every other operation of the form x @ W, it is per-token and requires no communication.

Figure 8: A routing table for four tokens under top-2 routing. Column e0 is already overloaded relative to e2 and e3, as a consequence of the router’s learned scores alone.

Load imbalance therefore originates in the model, not in the system. The routing table above is skewed before any question of device placement arises.

4.2 Dispatch and combine

Expert weights are large and are not moved. The tokens are moved instead. If a token’s selected experts reside on other devices, its hidden state must travel to them, and this transfer is called dispatch. Each expert then runs its own FFN over whatever tokens arrived, grouped by expert so that one matrix multiplication, a grouped GEMM, serves all tokens routed to it. The results travel back to the tokens’ origins, which is called combine, and are summed there.

Figure 9: Dispatch and combine for a single token routed to two experts on different devices.

Combine is simpler than the context-parallel merge of Section 3.4. That merge required log-sum-exp rescaling because a softmax normalizer cannot be computed correctly from a partial view of the keys. The MoE gate has already computed a full softmax over the selected experts locally, before dispatch, so the weights are correct on arrival and combine is a plain weighted sum.

4.3 Stragglers

Dispatch and combine are not per-token operations. Every token in the current batch is packed into a single collective over the expert-parallel group, and that collective does not complete for any participant until it has completed for all of them.

Figure 10: Token counts per rank, in parentheses, translate directly into transfer time. The lightly loaded ranks idle at the barrier until the heaviest one completes.

Consider a rank hosting an expert that the router favors heavily in the current batch. That rank pays for the same imbalance three times in succession, at three separate cost centers, each scaling with the same token count:

1. Dispatch receive. More tokens routed to its experts means more bytes arriving over the interconnect.
2. Compute. The grouped GEMM over those tokens is proportionally larger.
3. Combine send. The results, one vector per token, travel back out.

The lightly loaded ranks idle twice: once waiting for the heavy rank’s dispatch to complete, and again waiting for its compute and combine. Their hardware is free during both intervals, but the next layer cannot begin, because it requires every rank’s tokens to be assembled first.

This is the second exception announced in Section 1, and the target of the expert load balancing literature. EPLB replicates hot experts on periodically recomputed placements. UltraEP recomputes the replication plan from the exact post-gating load, on the critical path, every layer.

Continuous batching addresses a different problem. It governs which requests are admitted to or retired from the batch between iterations. Within a single iteration, every admitted token is packed into the same kernels and the same collectives, and stragglers arise at that level.

5. Collective communication

Both exceptions reduce to a small set of named collective operations. This section defines them, gives their costs, and identifies which ones the preceding sections have been invoking.

5.1 The primitives

Table 2: The standard collective operations, for P ranks.

Broadcast and scatter share a fan-out shape and differ in what travels: identical copies in the first case, distinct pieces in the second. Gather, all-gather, reduce, and all-reduce share a converge shape and differ along two dimensions: whether the center performs a computation (reduce) or concatenates (gather), and whether the result is delivered to one rank or to all.

Figure 11: Reduce and all-reduce differ only in whether the result is returned to one rank or to all of them.

All-reduce is the operation that combines the two partial outputs in tensor parallelism (Section 3.2). Reduce-scatter shares the converge step but distributes slices of the result instead of copies of it. In practice all-reduce is not implemented as a distinct primitive: it is a reduce-scatter followed by an all-gather.

All-to-all does not fit the converge-diverge shape. Every rank sends distinct data to every other rank simultaneously, and the operation is a transpose of the block matrix whose entry (i, j) is the payload rank i owes rank j.

Figure 12: All-to-all is a transpose of the send matrix. MoE dispatch is this operation, with token payloads in place of the indices shown.

5.2 The cost model

The standard model decomposes the time of a communication step into two terms that behave differently:

time  ≈  (number of messages) · α  +  (bytes moved) · β

where α is a fixed per-message latency, incurred regardless of payload size, and β is the inverse of bandwidth, the cost per byte. An operation is latency-bound when the first term dominates and bandwidth-bound when the second does. Which regime an operation falls into depends on the algorithm used to implement it, and not only on the collective chosen.

5.3 Ring all-reduce

All-reduce shows how much the algorithm matters. Two implementations produce identical results at very different cost.

A naive implementation reduces onto a root rank and broadcasts the result back. The root receives a full n-byte contribution from each of the other P - 1 ranks, so the busiest rank moves (P - 1) · n bytes, growing linearly in P.

The ring algorithm passes data around a cycle. It proceeds in two phases, a reduce-scatter followed by an all-gather, each costing (P - 1)/P · n bytes per rank, for a total of

2 (P - 1) / P · n

which approaches 2n from below as P grows and is effectively independent of the number of ranks.

Figure 13: Bytes moved per rank under naive and ring all-reduce, as the group size grows.

Ring all-reduce is the default in most distributed training frameworks for this reason. The ring requires 2(P - 1) sequential steps, so the accumulated α term can dominate for small messages, where a tree-based algorithm with O(log P) depth is faster. Production implementations such as NCCL select an algorithm based on message size.

5.4 Cost summary

Table 3: Bandwidth cost per rank under bandwidth-optimal algorithms. Here n denotes the size of the full buffer being reduced or gathered.

All-gather is the one operation whose cost grows with P unavoidably, since the combined result is P times a single rank’s contribution and no algorithm can avoid delivering that much new information.

All-to-all is bandwidth-comparable to reduce-scatter, but it requires P - 1 separate messages. In MoE dispatch many of those messages carry small payloads, since a given rank may route only a few tokens to a given distant expert. The fixed α term is then charged P - 1 times against a small β term, and the operation becomes latency-bound.

NanoCP’s routing-based communication backend targets this term. A standard collective imposes the full P × P message structure regardless of how much real traffic each pair carries. A routing table that issues transfers only for the pairs with genuine payloads reduces the message count directly, which a general-purpose collective library cannot do on the application’s behalf.

6. Conclusion

The two systems that motivated this article each attack one of the exceptions identified in Section 1.

NanoCP attacks the sequence-axis exception. A uniform context-parallel degree forces short requests to pay for a cross-device attention exchange they do not need, and binding a request’s KV cache to the same instance that performs its MoE dispatch makes it impossible to balance attention load and dispatch load at the same time. NanoCP decouples those two bindings and sizes the context-parallel degree per request.

UltraEP attacks the dispatch exception. Expert popularity is non-stationary at the granularity of a single microbatch, so any placement computed from historical statistics is already stale when it is used. UltraEP solves a replication and rerouting plan from the exact post-gating load, on the critical path, every layer.

Both systems reduce to decisions about which collective to invoke, over which group, and how sparsely. The cost model of Section 5 prices those decisions. NanoCP’s routing backend reduces the message count of the all-to-all, and UltraEP’s replication reduces the imbalance that makes the slowest rank, rather than the average rank, determine the cost of the collective.

Appendix A: A reference implementation

The following is a complete forward pass, from token identifiers to next-token logits, through two transformer layers, using the toy dimensions carried through this article (d_model = 8, n_heads = 2, T = 4, d_ff = 16). It uses NumPy only, and prints the shape and value of every intermediate quantity, including the masked score matrix of Section 2.3.

import numpy as np

np.random.seed(42)
np.set_printoptions(precision=2, suppress=True)

vocab_size = 6
d_model    = 8
n_heads    = 2
d_head     = d_model // n_heads
d_ff       = 16
T          = 4      # sequence length
num_layers = 2

def section(title):
    print("\n" + "=" * 78)
    print(title)
    print("=" * 78)

def layernorm(x, eps=1e-5):
    mu  = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps)

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def gelu(x):
    return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)))

# causal mask: query position i may only see key position j <= i
causal_mask = np.triu(np.ones((T, T)), k=1).astype(bool)   # True = masked out

# ---------------------------------------------------------------------
# STAGE 0: token ids -> embeddings -> + positional embedding
# ---------------------------------------------------------------------
section("STAGE 0: token ids -> embedding -> + positional embedding")

token_ids = np.array([2, 0, 4, 1])          # toy input sequence, length T=4
print("token_ids:", token_ids, " shape:", token_ids.shape)

embed_table = np.random.randn(vocab_size, d_model) * 0.5
x = embed_table[token_ids]                  # (T, d_model)
print("\nx = embedding lookup, shape:", x.shape)
print(x)

pos_embed = np.random.randn(T, d_model) * 0.1
x = x + pos_embed
print("\nx after + positional embedding, shape:", x.shape)
print(x)

# ---------------------------------------------------------------------
# Transformer layers
# ---------------------------------------------------------------------
for layer in range(num_layers):

    section(f"LAYER {layer} -- pre-attention LayerNorm")
    x_ln = layernorm(x)
    print("x_ln shape:", x_ln.shape)
    print(x_ln)

    section(f"LAYER {layer} -- project to Q, K, V  (x_ln @ W_Q / W_K / W_V)")
    W_Q = np.random.randn(d_model, d_model) * 0.3
    W_K = np.random.randn(d_model, d_model) * 0.3
    W_V = np.random.randn(d_model, d_model) * 0.3
    W_O = np.random.randn(d_model, d_model) * 0.3

    Q = x_ln @ W_Q     # (T, d_model)
    K = x_ln @ W_K
    V = x_ln @ W_V
    print("Q shape:", Q.shape, "\n", Q)
    print("\nK shape:", K.shape, "\n", K)
    print("\nV shape:", V.shape, "\n", V)

    section(f"LAYER {layer} -- reshape last dim into (n_heads, d_head)")
    Qh = Q.reshape(T, n_heads, d_head).transpose(1, 0, 2)   # (n_heads, T, d_head)
    Kh = K.reshape(T, n_heads, d_head).transpose(1, 0, 2)
    Vh = V.reshape(T, n_heads, d_head).transpose(1, 0, 2)
    print("Qh shape:", Qh.shape, " = (n_heads, T, d_head)")
    print("Head 0 Q (T,d_head):\n", Qh[0])
    print("Head 1 Q (T,d_head):\n", Qh[1])

    section(f"LAYER {layer} -- Q @ Kᵀ -> scores (T,T), causal masked, softmax")
    head_outputs = []
    for h in range(n_heads):
        scores = (Qh[h] @ Kh[h].T) / np.sqrt(d_head)        # (T, T)
        scores_masked = np.where(causal_mask, -1e9, scores)
        print(f"\n-- Head {h} --")
        print("raw scores (T,T):\n", scores)
        print("causal-masked scores:\n", scores_masked)

        weights = softmax(scores_masked, axis=-1)
        print("softmax weights (each row sums to 1):\n", weights)
        print("row sums:", weights.sum(axis=-1))

        out_h = weights @ Vh[h]                              # (T, d_head)
        print(f"head {h} output = weights @ V, shape {out_h.shape}:\n", out_h)
        head_outputs.append(out_h)

    section(f"LAYER {layer} -- concat heads, x W_O, residual add")
    concat = np.concatenate(head_outputs, axis=-1)           # (T, d_model)
    print("concat heads shape:", concat.shape, "\n", concat)
    attn_out = concat @ W_O
    print("\nattn_out = concat @ W_O, shape:", attn_out.shape, "\n", attn_out)

    x = x + attn_out
    print("\nx after residual add, shape:", x.shape, "\n", x)

    section(f"LAYER {layer} -- pre-FFN LayerNorm -> FFN (up, GELU, down) -> residual")
    x_ln2 = layernorm(x)
    W_up   = np.random.randn(d_model, d_ff) * 0.3
    W_down = np.random.randn(d_ff, d_model) * 0.3

    hidden = gelu(x_ln2 @ W_up)
    print("hidden = gelu(x_ln2 @ W_up), shape:", hidden.shape, "\n", hidden)

    ffn_out = hidden @ W_down
    print("\nffn_out = hidden @ W_down, shape:", ffn_out.shape, "\n", ffn_out)

    x = x + ffn_out
    print("\nx after FFN residual add, shape:", x.shape, "\n", x)

# ---------------------------------------------------------------------
# Final: LayerNorm -> unembed -> logits -> next-token probabilities
# ---------------------------------------------------------------------
section("FINAL: LayerNorm -> unembed -> logits -> next-token probabilities")

x_final = layernorm(x)
W_unembed = np.random.randn(d_model, vocab_size) * 0.3
logits = x_final @ W_unembed                                 # (T, vocab_size)
print("logits shape:", logits.shape, "\n", logits)

probs = softmax(logits, axis=-1)
print("\nnext-token probabilities per position (rows sum to 1):\n", probs)
print("row sums:", probs.sum(axis=-1))

next_token_pred = probs.argmax(axis=-1)
print("\nargmax predicted token id at every position:", next_token_pred)
print("only the LAST position's row is the actual 'next token' prediction:",
      next_token_pred[-1])

Two details in the output correspond to Section 2.3. Row 0 of the softmax weights collapses to [1, 0, 0, 0], since the first token can attend only to itself under the causal mask. Row 3 is the only row with a nonzero weight in all four columns, since the last query is the only one permitted to see the entire sequence.

Appendix B: What the feed-forward network computes

Section 2.5 gives the structural facts about the FFN that the main argument requires. This appendix covers what the component computes, which is a separate question and one the interpretability literature has answered in some detail.

B.1 The key-value memory interpretation

Geva et al. (EMNLP 2021) show that each of the d_ff inner units behaves like one entry of a learned associative memory. The columns of W_up act as keys, pattern detectors that a token’s vector is compared against by dot product. The activation function determines how strongly each pattern fires. The rows of W_down act as values, the vectors added to the residual stream when the corresponding pattern fires. The authors report that the learned patterns are human-interpretable, with lower layers capturing shallower patterns and upper layers capturing more semantic ones.

Figure 14: The FFN as a bank of learned pattern detectors. The labels are illustrative. No unit is designed to detect a given concept.

This interpretation applies only to trained weights. In the reference implementation of Appendix A, W_up and W_down are random and the inner units mean nothing. The claim is that gradient descent discovers such detectors, because they are an effective way of reducing next-token prediction loss.

B.2 What the value vectors do

A follow-up (Geva et al., EMNLP 2022) makes the function of the value vectors concrete. Projecting an individual value vector through the model’s unembedding matrix, the same matrix that converts a final hidden state into vocabulary logits, shows that it promotes a coherent cluster of related tokens. One value vector may raise the probability of a group of geographically related words, another the probability of a group of arithmetic-related words.

The FFN’s contribution to the residual stream is therefore a set of additive, direction-specific votes over the output vocabulary, weighted by how strongly each corresponding pattern matched the token. The authors demonstrate the causal force of this reading by suppressing value vectors associated with undesired token clusters and observing the corresponding drop in those outputs.

B.3 The residual stream

Each layer reads from the residual stream and writes back by addition. Elhage et al. (2021) describe the residual stream as a communication channel between layers, into which each block writes a linear projection of its output while leaving prior contributions intact.

Given the reading in B.2, addition preserves the accumulated votes of every previous layer. Under replacement, layer 40 could not build on anything layer 3 established, and intermediate representations would carry no interpretable signal. The additivity is also what makes the logit lens possible: applying the unembedding matrix to an intermediate residual state yields a meaningful, if rough, prediction, because that state already contains the summed contributions of all preceding layers.

B.4 Relation to attention

Geva et al. write the FFN formula to mirror attention’s, and the two follow the same match, weight, blend template:

attention:  output_i = Σ_j  softmax(Q_i · K_j)      · V_j
FFN:        output_i = Σ_c  GELU(x_i · W_up[:, c])  · W_down[c, :]

Two differences matter. Attention’s keys are dynamic: K = x @ W_K is recomputed from the surrounding context on every forward pass. The FFN’s keys are static, fixed at the end of training. This is why the FFN requires no cross-token communication while attention’s score step does, and it restates the criterion of Section 2.4 in a second setting. Attention’s weights are also normalized by a softmax and must sum to one across positions, where the FFN’s activations are unconstrained: all d_ff units may fire strongly, or none may.

References

The two papers this article was written alongside:

• Chen, J. et al. “NanoCP: Request-Level Dynamic Context Parallelism for Data-Expert Parallel Decoding.” arXiv:2605.21100
• Wei, X. et al. “UltraEP: Unleash MoE Training and Inference on Rack-Scale Nodes with Near-Optimal Load Balancing.” arXiv:2606.04101

Attention, tensor parallelism, and MLA:

• Shoeybi, M. et al. “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.” arXiv:2104.04473
• Brenndoerfer, M. “Tensor Parallelism: Column, Row, and Megatron Patterns.” mbrenndoerfer.com
• Jang, I. “Analyzing Parallelization of Attention.” insujang.github.io
• DeepSeek-AI. “DeepSeek-V2 Technical Report.” arXiv:2405.04434
• Vizuara. “Decoding Multi-Head Latent Attention,” Parts 1 and 2. vizuara.substack.com
• Towards Data Science. “DeepSeek-V3 Explained 1: Multi-head Latent Attention.” towardsdatascience.com
• Bhatia, N. et al. “Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding.” arXiv:2507.07120
• vLLM GitHub. RFC for Helix Parallelism implementation. github.com/vllm-project/vllm/issues/34018
• ROCm. “The vLLM MoE Playbook.” rocm.blogs.amd.com
• vLLM Docs. “Data Parallel Deployment.” docs.vllm.ai
• Jarvislabs. “Scaling LLM Inference: DP, PP, TP.” docs.jarvislabs.ai

Context parallelism and the online-softmax merge:

• Dao, T. et al. “Flash-Decoding for long-context inference.” pytorch.org/blog
• Meta AI. “Context Parallelism for Scalable Million-Token Inference.” arXiv:2411.01783
• HuggingFace Blog. “FlashAttention: Online Softmax.” huggingface.co/blog

Serving architecture, memory, and scheduling:

• Kwon, W. et al. “Efficient Memory Management for Large Language Model Serving with PagedAttention,” explained by Brenndoerfer, M. mbrenndoerfer.com
• Zhong, Y. et al. “DistServe: Disaggregating Prefill and Decoding for Goodput-Optimized Large Language Model Serving.” arXiv:2401.09670
• Agrawal, A. et al. “Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve.” OSDI 2024.
• NVIDIA. “GB200 NVL72.” nvidia.com
• Introl. “NVLink and Scale-Up Networking.” introl.com

Stragglers and MoE communication:

• Ho, Q. et al. “Solving the Straggler Problem for Iterative Convergent Parallel ML.” CMU-PDL Technical Report. pdl.cmu.edu
• Zuo, P. et al. “Serving Large Language Models on Huawei CloudMatrix384.” arXiv:2508.02520
• DeepSeek-AI. “DeepEP: an efficient expert-parallel communication library.” github.com/deepseek-ai/DeepEP
• DeepSeek-AI. “EPLB: Expert Parallelism Load Balancer.” github.com/deepseek-ai/EPLB
• ROCm. “Dropless MoE Training with Primus-Turbo.” rocm.blogs.amd.com

FFN interpretability and the residual stream:

• Geva, M., Schuster, R., Berant, J., Levy, O. “Transformer Feed-Forward Layers Are Key-Value Memories.” EMNLP 2021. aclanthology.org
• Geva, M., Caciularu, A., Wang, K., Goldberg, Y. “Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space.” EMNLP 2022. aclanthology.org
• Elhage, N. et al. “A Mathematical Framework for Transformer Circuits.” Transformer Circuits Thread, 2021. transformer-circuits.pub
• “MLPs in Transformers.” Learn Mechanistic Interpretability. learnmechinterp.com