Attention, Parallelism, and MoE Communication: A Ground-Up Walkthrough

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.

Part 1: Inside one attention block

What is W_K, really? Where does it sit in the pipeline?

A token’s “hidden state” going into an attention block is a vector of length d_model. A batch of B sequences, each T tokens long, gives a tensor of shape (B, T, d_model). W_K is just a plain, non-head-aware (d_model, d_model) matrix. Multiplying x @ W_K gives back a flat (B, T, d_model) tensor with no visible head structure yet. Q and V are produced the same way, each with their own separate weight matrix.

The only step where token positions actually interact is the amber box: Q @ Kᵀ. Everything else (projections, softmax normalization, output projection) is per-token or per-head local. That single fact turns out to explain almost everything else in this post.

How does “splitting heads” actually work mechanically?

Reshaping d_model = n_heads × d_head is pure relabeling, no computation. Take one token’s 8-number K-vector: chopping it into two groups of 4 consecutive numbers is “head 0” and “head 1.” Nothing forces that split point; it’s just a convention that has to be applied consistently to Q, K, and V so head 0’s pieces line up across all three.

Because the split happens on the output columns of the weight matrix itself, you can literally hand GPU 0 the left half of W_K’s columns and GPU 1 the right half. Each GPU computes its own head from scratch, with zero communication needed until the very end.

For TP’s combine step: each GPU multiplies its head output by its own row-slice of W_O (not the full matrix), producing a partial d_model-sized output, and the true answer is the sum of those partials, not a concatenation. So the collective that combines TP’s two GPUs is an all-reduce, not an all-gather: a distinction that matters later.

Why doesn’t MLA (DeepSeek’s attention variant) shard cleanly this way?

Standard multi-head attention has a natural per-head cache: head 1’s K/V lives in its own slice. MLA throws that away: it compresses K and V into one shared latent vector per token, and every head reconstructs its own K/V from that same shared vector via a per-head up-projection. That compression is the entire reason MLA’s KV cache is so small.

The problem: if you TP-shard by head, every GPU still needs the entire shared latent to reconstruct even its own heads. There’s no per-head slice to hand out, because the compression collapsed it into one shared thing.

This is why modern MoE serving stacks (and both papers this post is built around) run DP-EP, not TP-EP: DP replicates attention entirely (no sharing, no duplication problem), and EP is a completely different axis (sharding experts, not KV cache) used for the FFN.

Part 2: DP, TP, CP, PP: which axis does each one cut?

Once Q/K/V are reshaped to (B, T, n_heads, d_head), every axis of that tensor is a candidate for a different parallelism strategy:

The rule that decides all four: an axis needs communication only when the computation must combine information across different slices of it. B and n_heads are axes attention treats as fully independent (batch elements never interact; heads never interact until the very last step). T is the axis attention is defined to reduce over: a query needs to see keys from potentially every position, so cutting T genuinely separates data the computation needs together. d_head would split a single dot product mid-sum, essentially never attempted.

Ranking these by actual communication cost gives B < n_heads < T < d_head, not simply “deeper in the tensor shape = more communication.” n_heads (TP) needs one cheap combine step at the very end; T (CP) needs genuine ongoing exchange, which is exactly why CP is a whole research subfield and TP mostly isn’t.

How does CP’s communication actually work?

Splitting T for the Q/K/V projection is exactly as free as DP splitting B: the projection is per-token, so a GPU can compute Q/K/V for its own local tokens with zero communication. The break happens one step later, at the score computation, because that’s the one operation whose entire job is “let every token look at every other token”:

Two real implementations exist for that cross-boundary exchange. NanoCP/Helix route the query directly to whichever remote GPU holds the needed KV shard, get a partial result + a scaling factor back, and merge (the same online-softmax trick FlashAttention uses internally). Ring Attention does it differently: no fixed destination, just a rotating relay:

Ring trades a longer, bandwidth-friendly relay for training-scale CP groups; the routed/point-to-point style is lower-latency for the smaller, more dynamic groups decode-time serving needs, which is exactly the design choice NanoCP makes.

On PP: worth a note since it’s easy to assume it’s training-only. It’s genuinely preferred over TP once GPUs span more than one fast-interconnect domain, because TP needs an all-reduce every layer (brutal over slow links) while PP only passes activations across one boundary per stage. Sarathi-Serve reports roughly 2× lower median latency for PP vs. TP once crossing nodes over ordinary Ethernet. This is exactly why UltraEP’s DeepSeek-V3 config uses EP64-PP4 rather than adding TP.

Part 3: The toy transformer

Rather than keep everything abstract, here’s a complete, runnable forward pass, embeddings through two full transformer layers to next-token logits, with the same toy dimensions used throughout (d_model=8, n_heads=2, T=4). Every stage prints its actual shape and values.

import numpy as np

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

# ---------------------------------------------------------------------
# Toy dimensions -- same numbers used throughout the conversation
# ---------------------------------------------------------------------
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])

The masked score matrix is the part worth staring at: row 0 collapses to [1,0,0,0] after softmax (the first token can only attend to itself), while row 3 gets a genuine distribution across all four positions: proof the causal mask is doing exactly what it should.

Part 4: What is the FFN actually for?

What do W_up and W_down do, intuitively?

Shape-wise: W_up is (d_model, d_ff): expansion. W_down is (d_ff, d_model): contraction back to fit the residual add. But the functional answer comes from real interpretability research (Geva et al., EMNLP 2021, “Transformer Feed-Forward Layers Are Key-Value Memories”): each of the d_ff hidden units behaves like one entry in a lookup table. W_up’s columns are keys: learned pattern detectors a token gets dot-producted against. GELU is the gate deciding how strongly each pattern fired. W_down’s rows are values: what gets added if that pattern fired.

I proved the “collapse without GELU” claim rather than just asserting it: (x @ W_up) @ W_down computed directly is numerically identical to x @ (W_up @ W_down) collapsed into a single (d_model, d_model) matrix, and np.allclose returns True. Add GELU back in and the two stop matching entirely. The nonlinearity is the entire reason the wide middle layer buys any extra capacity.

What does the “value” actually promote?

This is the part that made it click for me. A follow-up paper (Geva et al., EMNLP 2022) projects each value-vector through the model’s unembedding matrix, the same matrix used at the very end to turn a hidden state into vocabulary logits, and finds each one promotes a coherent cluster of actual words:

“A value vector might promote the cluster {Paris, France, French, European, Seine}, or the cluster {multiply, divide, arithmetic, calculate}.”

Important honesty check: this story only exists after training on real data. In the toy script above, W_up/W_down are random noise. There are no real “capital cities” detectors in there. The finding is that gradient descent discovers useful pattern-detectors on its own during training, purely because they’re a useful building block for predicting the next token well.

Why add to the residual stream instead of overwriting it?

Because the value vectors are literally votes toward specific output words, and addition is the only operation that lets those votes accumulate instead of erasing each other. This is Anthropic’s own framing from “A Mathematical Framework for Transformer Circuits” (Elhage et al., 2021):

“Rather than overwriting the results from previous layers, each layer of the residual block ‘reads’ its input from the residual stream… and then ‘writes’ its result to the residual stream by adding a linear projection back in.”

If FFN layers overwrote instead of added, deep networks would be pointless: layer 40 couldn’t build on anything layer 3 discovered.

Is W_down doing something similar to V in Q/K/V?

Yes, genuinely, not just an analogy. Geva et al. deliberately wrote the FFN formula to mirror attention’s: match → weight → blend, in both cases. Attention: output = Σ softmax(Q·K) × V. FFN: output = Σ GELU(x·W_up) × W_down. Same template.

The one real, load-bearing difference: attention’s keys are dynamic: K = x @ W_K, computed fresh from whatever’s in context. The FFN’s keys are static, learned parameters, fixed the moment training ends. That’s precisely why the FFN needs zero cross-token communication while attention’s score step is the one place that genuinely does. (Second difference, more of a footnote: attention’s weights are forced to sum to 1 via softmax; GELU has no such constraint, all d_ff keys can fire independently, unconstrained.)

Part 5: MoE: dispatch and combine

MoE replaces the single shared FFN with many smaller, specialized ones (“experts”), plus a gate that decides, per token, which top-k experts should process it. The gate itself is a per-token linear layer, no communication needed yet.

If e0 and e1 live on one GPU while e2/e3 live on another, each token’s hidden state physically has to travel to wherever its chosen experts are (dispatch), get processed there, and travel back to be blended (combine):

Combine here is simpler than CP’s merge: the gate already normalized w0, w1 locally before dispatch, so it’s a plain weighted sum, no rescaling trick required.

Why does this create a straggler?

All 4 tokens’ dispatch calls are packed into one shared collective, not 4 independent sends. If e0 gets picked by every token and e1 gets picked by fewer, the GPU hosting e0 ends up with far more total traffic than the others, and the collective as a whole only completes once the busiest GPU finishes:

The idle GPUs’ hardware is fully free during the dashed span. There’s just nothing useful for them to do, because the next layer’s computation needs every GPU’s dispatched tokens assembled together before it can proceed at all. This is exactly the problem UltraEP solves by replicating a hot expert’s weights onto a spare GPU in real time and splitting the token load across the replicas, reacting to the exact, current-iteration load rather than periodic, stale statistics the way its predecessor (EPLB) does.

One clarification worth keeping: continuous batching solves a different problem than this. It changes which requests are included between iterations. Within one iteration, every included token is glued into the same batched kernel calls and the same collective, that’s the level the straggler problem actually lives at.

Part 6: The collective operations underneath everything

Everything above eventually reduces to a handful of named collective operations. Worth knowing them by name and cost, not just by what we’ve called them informally.

Broadcast / Scatter: one rank’s data either copied identically to everyone (broadcast), or cut into pieces with one distinct piece per rank (scatter). Gather / All-Gather: the reverse; everyone’s data converges onto one rank (gather) or onto every rank (all-gather). Reduce / All-Reduce: same converge shape, but the center step does real work (a sum), landing on one rank or all of them:

All-reduce is exactly what combined TP’s two head outputs earlier. Reduce-scatter is the same converge step, but instead of broadcasting the whole sum back, each rank keeps only its own slice, and in real implementations, all-reduce isn’t a separate primitive at all, it’s reduce-scatter followed by all-gather, chained together.

All-to-all is the odd one out: no converge/diverge shape at all, just every rank sending something different to every other rank simultaneously. It’s literally a matrix transpose:

How much does each of these actually cost?

The standard tool is the α-β model: total time ≈ (number of messages) × α (fixed latency per message) + (bytes moved) × β (per-byte cost). Whether an operation is latency-bound or bandwidth-bound depends on which term dominates, and that’s where the choice of algorithm, not just which collective, matters enormously.

The single most important concrete result here: naive all-reduce (everyone sends to rank 0, it sums, it broadcasts back) costs (P-1) × n bytes on the busiest rank, linear in the number of ranks. Ring all-reduce (pass data around a ring instead) costs 2(P-1)/P × n, which approaches 2n and flatlines as P grows, essentially independent of rank count:

That flat curve is why ring all-reduce became the default in every serious distributed training framework. The tradeoff: ring needs 2(P-1) sequential steps, so for tiny messages the fixed per-message latency can dominate anyway. Real implementations (NCCL) switch strategy based on message size.

The rest, briefly: broadcast / scatter / gather move ~n bytes per rank, roughly P-independent. All-gather genuinely does scale with P: the combined result is P·n, so each rank must receive (P-1)·n new bytes; no algorithm avoids this, since the information itself grows. Reduce-scatter is exactly half of ring all-reduce (the first phase). All-to-all is bandwidth-comparable to reduce-scatter, but its real cost is latency: it needs P-1 distinct messages, and if each one carries only a handful of tokens (the common MoE case), the fixed per-message overhead α dominates over the actual bytes. That’s the precise reason NanoCP builds a routing-based backend instead of calling a generic all-to-all: it attacks the message-count term directly, skipping pairs that have nothing real to send, rather than paying for a dense P×P mesh every single request.

References

The two papers this whole post is built around:

• 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:

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

Context parallelism, FlashAttention, and the online-softmax merge:

• Dao, T. et al. “Flash-Decoding for long-context inference.” pytorch.org/blog (mirrored at princeton-nlp.github.io)

• 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

• NVIDIA, “GB200 NVL72.” nvidia.com

• Introl, “NVLink and Scale-Up Networking.” introl.com

Straggler dynamics 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
• “MLPs in Transformers.” Learn Mechanistic Interpretability. learnmechinterp.com
• Elhage, N. et al. “A Mathematical Framework for Transformer Circuits.” Transformer Circuits Thread, 2021. transformer-circuits.pub

Pipeline parallelism in serving:

• Agrawal, A. et al. “Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve.” OSDI 2024.
• “RServe” (pipeline-parallel serving latency comparison against tensor-parallel vLLM).