Tensor-Parallel Inference

"They cut me into eight slices and told me I would feel faster. They were right, until layer eighty-one, when I spent more time gossiping than thinking."

A Weight Matrix Sharded One Slice Too Far

In the previous section we saw why a large language model can outgrow one GPU on two separate fronts: its weights may not fit, and even when they fit, a single GPU may compute each token too slowly for the latency budget. Replication, the subject of Chapter 23, answers neither problem, because a replica is a full copy of the model on one device; if the model does not fit on one device, you cannot make a replica at all, and adding replicas never makes a single request faster. The remedy is to make one logical model span several GPUs. Tensor parallelism is the first and most latency-friendly way to do that, and it is, almost beat for beat, the inference reuse of a method this book already developed for training.

1. The Same Split, Now Forward-Only Intermediate

Recall the Megatron-style split from Section 16.2. A transformer layer is, at its core, two large matrix multiplications back to back: an up-projection that widens the hidden vector, a nonlinearity, and a down-projection that narrows it again. Tensor parallelism partitions those matrices along their shared inner dimension. The first weight matrix is split by columns, so each of the $T$ GPUs holds a vertical slice and computes only a slice of the widened activation. The second weight matrix is split by rows to match, so each GPU multiplies its activation slice by its own row slice and produces a partial copy of the full-width output. No GPU holds the whole layer; each holds a $1/T$ slice of both weight matrices, and the attention block splits the same way, by attention heads.

The partial outputs are the key. Each GPU computes a full-shape output vector that is correct except that it is missing the contributions of the slices held by the other GPUs. Because the down-projection is a sum over the inner dimension, and a sum decomposes exactly the way the gradient sum did in Section 1.1, adding the $T$ partial outputs reconstructs the exact single-GPU result. That addition is an all-reduce: every GPU contributes one vector, the vectors are summed, and every GPU receives the total so it can feed the next layer. This is the same all-reduce collective introduced in Section 4.3, now fired once on the forward pass of each layer rather than once on the backward pass of each step. Figure 24.2.1 lays out this split across three GPUs, with the column-and-row weight slices, the head-sharded KV cache, and the per-layer all-reduce that combines the partial outputs.

Two consequences follow immediately, and they are exactly why tensor parallelism is used at inference rather than only at training. First, the weight memory per GPU drops by a factor of $T$, because each GPU stores only its slice; a model that needs four GPUs' worth of memory becomes servable on four GPUs that no single one could host. Second, the per-token latency drops, because the two big matrix multiplications that dominate each layer are now done $T$ ways in parallel, with each GPU doing $1/T$ of the arithmetic. Replication can give you neither: it cannot fit a model that does not fit, and it cannot make one token faster. Tensor parallelism makes ONE replica span GPUs; replication, the contrast we draw in Section 5, adds MORE replicas. They are orthogonal, and a real serving stack uses both.

2. Verifying the Split Equals One GPU Intermediate

The claim that summing partial outputs reconstructs the single-GPU result is the whole correctness argument, and it is worth seeing run rather than asserted. The code below builds one transformer feed-forward block, computes its output the ordinary single-GPU way, then computes it again as $T$ sharded GPUs that each hold a column slice of the first weight, the matching row slice of the second weight, and contribute one partial output to an all-reduce. It checks that the two agree for several values of $T$.

import numpy as np

rng = np.random.default_rng(7)
d_model, d_ff, S = 256, 1024, 16      # hidden dim, FFN inner dim, sequence length

x  = rng.standard_normal((S, d_model)) / np.sqrt(d_model)
W1 = rng.standard_normal((d_model, d_ff)) / np.sqrt(d_model)
W2 = rng.standard_normal((d_ff, d_model)) / np.sqrt(d_ff)

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

# Single-GPU reference: y = gelu(x @ W1) @ W2
ref = gelu(x @ W1) @ W2

# Tensor-parallel across T simulated GPUs.
# W1 column-parallel (slice the d_ff axis), W2 row-parallel (same axis).
# Each GPU returns a full-shape partial; summing them is the all-reduce.
def tensor_parallel(T):
    cols = np.array_split(np.arange(d_ff), T)        # how the inner dim is sliced
    partials = []
    for c in cols:                                    # each c lives on one GPU
        W1_shard = W1[:, c]                            # column slice of W1
        W2_shard = W2[c, :]                            # matching row slice of W2
        h = gelu(x @ W1_shard)                         # local activation (S, |c|)
        partials.append(h @ W2_shard)                 # partial (S, d_model) output
    return np.sum(partials, axis=0)                   # all-reduce SUM across GPUs

for T in (1, 2, 4, 8):
    out = tensor_parallel(T)
    print(f"T={T}  max|TP - single| = {np.max(np.abs(out - ref)):.2e}")

T=1  max|TP - single| = 0.00e+00
T=2  max|TP - single| = 2.50e-16
T=4  max|TP - single| = 2.64e-16
T=8  max|TP - single| = 2.50e-16

The agreement reported in Output 24.2.1 is at the level of double-precision rounding, so tensor parallelism is, like data parallelism before it, an exact reorganization and not an approximation. The reader who served a request expecting a degraded answer gets the identical answer. What changed is purely operational: the arithmetic that one GPU would have done alone is now spread across $T$ GPUs, with one all-reduce per layer stitching their partial outputs back into the single coherent token the model would have produced.

3. Why It Must Stay on the Fastest Interconnect Advanced

The all-reduce per layer is the entire reason tensor parallelism is confined, in practice, to a single node. A token-generation step walks every layer of the model in sequence, and each layer ends with an all-reduce that every GPU must finish before the next layer can begin. For an 80-layer model generating one token, that is 80 all-reduces standing directly on the critical path, with no useful computation to hide behind them, because the next layer literally depends on the summed result. This is the topology argument of Section 4.9 made concrete: a collective on the critical path must run on the fastest links available, and within a node those links are NVLink, which moves data between GPUs at hundreds of gigabytes per second, an order of magnitude faster than the network between nodes.

Push the tensor-parallel group across a node boundary and every one of those 80 per-token all-reduces now crosses the slow inter-node fabric. The collective that was a rounding error in the layer's time budget becomes the dominant term, and per-token latency rises instead of falls. This is why practical tensor-parallel degree is capped, usually to the number of GPUs in one node (8 on a typical server), and why spanning more machines is the job of pipeline parallelism in Section 24.3, which communicates far less often. Tensor parallelism is the within-node tool; the moment you leave the node, its economics invert.

4. The Latency-Versus-Overhead Trade-Off Advanced

If more GPUs always made each token faster, you would use as many as you could buy. They do not, and the reason is the all-reduce. We can model the per-token latency directly. Let the per-layer compute on one GPU cost $c_0$, falling as $c_0/T$ when split across $T$ GPUs. Let the per-layer all-reduce cost grow with $T$ as a ring or tree collective does, roughly $a + b\log_2 T$ for $T > 1$, where $a$ is a fixed latency floor and $b$ scales the per-step transfer term. Then the per-token latency for an $L$-layer model is

The first term falls with more GPUs; the second and third rise. Their sum has a minimum and then climbs, so adding GPUs helps until the all-reduce overhead overtakes the shrinking compute, after which more GPUs make each token slower, not faster. The same code that verified correctness also evaluates this model in Code 24.2.2, so the diminishing returns are measured rather than claimed.

# Latency model: compute falls as 1/T, all-reduce overhead grows with T.
# Units are milliseconds; values are illustrative of an NVLink-class node.
c0, a, b, L = 1.00, 0.04, 0.06, 80     # per-layer compute, allreduce floor, allreduce slope, layers
print(f"{'T':>3} {'compute':>9} {'allreduce':>10} {'total/layer':>12} {'token(ms)':>10} {'speedup':>8}")
base = None
for T in (1, 2, 4, 8, 16):
    compute = c0 / T
    allred  = (a + b * np.log2(T)) if T > 1 else 0.0   # no collective with one GPU
    per_layer = compute + allred
    token = per_layer * L                              # L layers on the critical path
    base = token if base is None else base
    print(f"{T:>3} {compute:>9.3f} {allred:>10.3f} {per_layer:>12.3f} {token:>10.2f} {base/token:>7.2f}x")

  T   compute  allreduce  total/layer  token(ms)  speedup
  1     1.000      0.000        1.000      80.00      1.00x
  2     0.500      0.100        0.600      48.00      1.67x
  4     0.250      0.160        0.410      32.80      2.44x
  8     0.125      0.220        0.345      27.60      2.90x
 16     0.062      0.280        0.342      27.40      2.92x

The numbers tell the practical story crisply. Two GPUs give a healthy $1.67\times$; eight give $2.9\times$, already well below the ideal $8\times$ because every halving of compute is partly eaten by a growing collective. Sixteen GPUs are essentially pointless here, gaining 0.02 of a multiplier while doubling the hardware bill, and on real hardware where 16 GPUs would straddle two nodes the curve would bend the other way and latency would rise. This is the quantitative form of the rule from Section 3: there is a sweet spot for the tensor-parallel degree, it is usually within one node, and past it you are paying for GPUs that make the model slower.

5. The Sharded KV Cache, and the Contrast with Replication Intermediate

Weights are not the only thing tensor parallelism shards. The KV cache, the per-token attention state whose per-node economics Chapter 22 developed as a prerequisite, is also split across the tensor-parallel GPUs. Because attention is partitioned by head, each GPU computes and stores the keys and values only for its own heads, so the cache for a long context is divided $T$ ways alongside the weights. This is the multiplication promised at the end of Chapter 22: per-node KV-cache pressure, already the dominant memory cost of long-context serving, is here relieved by the same factor $T$ that relieves the weight memory, which is part of why tensor parallelism is what makes very long contexts servable at all. No extra collective is needed for the cache itself; each GPU's heads stay on that GPU, and the per-layer all-reduce that combines the attention output already carries the cross-head coupling.

It helps to state plainly how this differs from replication, the subject of Chapter 23. Replication makes copies: each replica is a complete model on its own GPU (or its own tensor-parallel group), and adding replicas raises how many requests per second the fleet can serve. Tensor parallelism makes one model bigger than one GPU: it spans a single logical model across several GPUs to fit it and to speed up each token. The two compose cleanly and are the two axes of a serving deployment, often written together as a configuration like "tensor-parallel size 8, with 6 replicas." You reach for tensor parallelism when one model will not fit or one token is too slow; you reach for replication when one request stream is more than the fleet can keep up with. Confusing them is a common and expensive mistake: adding replicas will never fix a model that does not fit, and adding tensor-parallel GPUs past the sweet spot will never raise throughput.

from vllm import LLM, SamplingParams

# tensor_parallel_size=4 shards every layer's weights and the KV cache across
# 4 GPUs on this node, and inserts the per-layer all-reduce on NVLink.
llm = LLM(model="meta-llama/Llama-3.1-70B-Instruct", tensor_parallel_size=4)

out = llm.generate(["Explain tensor parallelism in one sentence."],
                   SamplingParams(max_tokens=64))
print(out[0].outputs[0].text)

6. When Tensor Parallelism Is the Right Tool Intermediate

Tensor parallelism is the right answer to two questions and the wrong answer to a third. It is right when a model does not fit on one GPU, because it divides the weight and cache memory by $T$. It is right when a single token must be produced faster than one GPU can manage, because it divides the per-layer compute by $T$, up to the sweet spot the latency model exposed. It is the wrong answer when the model already fits and runs fast enough and you simply need more requests per second; that is replication's job, and reaching for tensor parallelism there just adds all-reduces to a model that did not need them. Keep the group inside one node, find the tensor-parallel degree at the bottom of the latency curve rather than the largest you can afford, and let pipeline parallelism (Section 24.3) carry the model across nodes when one node is genuinely not enough.

We now have the within-node tool for serving a model too large or too slow for one GPU: shard every layer's weights and KV cache across $T$ GPUs, pay one all-reduce per layer on NVLink, and stop at the degree where the collective overtakes the speedup. When one node is not enough, the model must be cut a different way, across nodes, with far less frequent communication. That is pipeline parallelism, and Section 24.3 takes it up.