In Megatron tensor parallelism, where do the column and row splits go in attention and the MLP, and why two all-reduces per layer?

TL;DR: Split the first matrix of each block by columns (QKV by heads, gate and up by intermediate width) and the second by rows (the output projection, the down projection), so the nonlinearity and the attention math stay local and each block ends with one all-reduce of a partial sum. Two blocks per layer means two all-reduces: for Llama 3.3 70B at TP=8 that is 160 per token and 4.59 MB sent per GPU per token.

How to approach it

Start from one block, $Y = f(XA)\,B$, and ask which split of $A$ lets $f$ run without talking to another GPU. Then apply the same pattern to attention, where the heads play the role of columns. Ask the TP degree and the model shape, because the head counts must divide by TP and the per-token bytes come from the hidden size. Mention the decode regime early: the payloads are small, so latency per call matters more than bandwidth.

A strong answer

The MLP. Split $A$ (gate and up, $d \times f$) by columns. GPU $r$ computes $XA_r$ for its slice of the intermediate width, and because SwiGLU is elementwise, $\mathrm{silu}(XA^{\text{gate}}_r) \odot XA^{\text{up}}_r$ needs nothing from other GPUs. Split $B$ (down, $f \times d$) by the matching rows. Each GPU then holds a full-width partial result, and the true output is their sum:

$$Y = \sum_{r=1}^{p} f(XA_r)\,B_r$$

That sum is the all-reduce. Reverse the order (row split first) and you would need a reduction before the nonlinearity, because $f(\sum_r Z_r) \neq \sum_r f(Z_r)$. That is the whole reason for column-then-row.

Attention. Heads are independent, so Q, K and V are split by columns along the head dimension. Each GPU owns whole heads, runs attention on them with its own slice of the KV cache, and the output projection $W_O$ is split by rows. Same pattern, same single all-reduce at the end.

The input to each block is replicated on every GPU, and the output after the all-reduce is replicated again, which is what the next block needs. So the count is one all-reduce per block and two per layer.

What it costs. Each all-reduce carries one hidden vector per token, 8,192 × 2 bytes in BF16. A ring all-reduce sends $2(p-1)/p$ of the buffer from each GPU:

payload per all-reduce per token = 8,192 x 2 B = 16,384 B
sent per GPU (ring, p = 8)       = 2 x 7/8 x 16,384 B = 28,672 B
per token, whole model           = 2 x 80 layers x 28,672 B = 4,587,520 B = 4.59 MB
batch 64 decode step             = 64 x 4.59 MB = 294 MB per GPU
at 450 GB/s per direction        = 294 MB / 450 GB/s = 0.65 ms
sanity: weight read per GPU = 141 GB / 8 = 17.6 GB, 5.27 ms at 3.35 TB/s

At batch 1 the payload is 16 KB, and 160 calls cost more in fixed latency than in bytes. With an assumed 10 µs per call that is 1.6 ms per token against a 5.27 ms weight floor, which is why TP stops scaling well past one node. The fixed cost per call is something to measure on the hardware; it is not a datasheet figure.

A reference check that the split is exact:

import numpy as np

def silu(x):
    return x / (1.0 + np.exp(-x))

def mlp_tp(x, w_gate, w_up, w_down, tp):
    """x: [T, d]; w_gate, w_up: [d, f]; w_down: [f, d]. Simulates tp ranks."""
    f = w_gate.shape[1]
    assert f % tp == 0, "intermediate size must divide by tp"
    s = f // tp
    partials = []
    for r in range(tp):                      # each loop body is one GPU
        cols = slice(r * s, (r + 1) * s)
        h = silu(x @ w_gate[:, cols]) * (x @ w_up[:, cols])  # column split: local
        partials.append(h @ w_down[cols, :])  # row split: partial sum [T, d]
    return np.sum(partials, axis=0)           # the one all-reduce

rng = np.random.default_rng(0)
d, f = 64, 256
x = rng.standard_normal((5, d))
wg, wu, wd = (rng.standard_normal(s) * 0.1 for s in [(d, f), (d, f), (f, d)])
ref = (silu(x @ wg) * (x @ wu)) @ wd
for tp in (1, 2, 8):
    assert np.allclose(mlp_tp(x, wg, wu, wd, tp), ref)
assert np.allclose(mlp_tp(x[:0], wg, wu, wd, 4), ref[:0])  # empty batch

When does the layout change? If the TP degree exceeds the KV head count (16 GPUs for 8 KV heads), K and V can no longer be split by head and get replicated, so the cache stops shrinking per GPU. And in prefill, where the all-reduce payload is thousands of tokens, engines often swap each all-reduce for a reduce-scatter plus all-gather (sequence parallelism) so norms and residuals run on a slice of the tokens. The bytes are the same; the memory and the elementwise work drop by $p$.

What interviewers probe next

"Why not split the down projection by columns too?" Its input would then have to be gathered in full on every GPU first, which is an all-gather of the 28,672-wide intermediate: 7/8 x 57,344 B = 50,176 B per GPU per token, 1.75x the hidden-size all-reduce, and the column-split output would then need gathering again.

"Where does the embedding and LM head go?" Vocabulary-parallel: each GPU holds a slice of the 128,256 rows, and the softmax needs a max and a sum reduced across GPUs, or a gather of the logits.

"What about TP=16 for this model?" The 8 KV heads cannot split 16 ways, so each KV head lives on two GPUs, and every all-reduce now crosses the node boundary unless the node has 16 GPUs in one NVLink domain.

How it's scored

Weak: says each GPU "takes a slice of the layers", confusing TP with pipeline parallelism, or cannot say why the first matrix is split by columns.

Hire: places all five splits, explains that the elementwise nonlinearity forces column-then-row, and counts two all-reduces per layer.

Strong Hire: prices the bytes per token with the ring factor, separates the latency-bound decode case from the bandwidth-bound prefill case, and names the KV-head limit on the TP degree.

Common mistakes

• Counting one all-reduce per layer, which misses that attention and the MLP each end in a row-split matrix.
• Splitting K and V across 64 query heads instead of the 8 KV heads, which suggests TP=64 is possible.
• Quoting the all-reduce bytes without the $2(p-1)/p$ factor, which understates them by 1.75x at TP=8.
• Using 900 GB/s as the per-direction NVLink figure for the transfer time.