Yesterday we used the roofline model to explain why LLM inference can be limited by memory traffic rather than arithmetic. We saw that batching improves weight reuse, while long contexts can make KV-cache reads the dominant cost.
Today we'll add another resource to that picture: communication between GPUs.
The intuition is:
Splitting a model across GPUs gives each GPU less work, but the GPUs must exchange intermediate results. Sometimes the conversation between them costs more than the computation it saves.
This is the central tradeoff in tensor parallelism (TP), a technique used in both LLM training and inference.
Why split one model across multiple GPUs?¶
A dense model with 70 billion BF16 parameters needs roughly
just for its weights. That excludes KV caches, temporary activations, and—in training—gradients and optimizer state.
One option is to divide the model's layers among GPUs. That's pipeline parallelism: GPU 1 runs early layers, then passes activations to GPU 2 for later layers.
Tensor parallelism does something different: multiple GPUs cooperate on the same layer by splitting its large weight matrices.
In an ideal four-way split, the 140 GB of weights becomes approximately 35 GB per GPU. Real allocations are not perfectly even because some parameters are replicated and other buffers still consume memory.
The key mathematical trick: split a matrix multiplication¶
Consider a linear layer:
where \(X\) contains token activations and \(W\) is the weight matrix.
Suppose we divide \(W\) by columns across two GPUs:
Then
Each GPU computes a different slice of the output. Neither needs the other GPU's answer yet:
input X
/ \
GPU 1 GPU 2
XW1 XW2
\ /
concatenate
YThis is column-parallel linear algebra. The output is distributed across devices; we may not need to concatenate it physically if the next operation can consume the shards directly.
Now consider a second linear layer whose input dimension is split across GPUs. Write
Its output is
Each GPU computes a partial answer, but those answers must be summed to obtain the full output. This is row-parallel linear algebra.
GPU 1: H1 × V1 = partial Y1 ──┐
├── SUM → full Y
GPU 2: H2 × V2 = partial Y2 ──┘The distributed sum is commonly implemented with an all-reduce collective: every participating GPU ends up with the same summed result. Some architectures use reduce-scatter instead and keep the result sharded.
Why Transformers are a natural fit¶
A Transformer block contains large projections in attention and in its feed-forward network (MLP).
A common tensor-parallel arrangement is:
- Attention: distribute query/attention heads across GPUs; combine their contributions through the output projection.
- MLP: split the expanding projections across GPUs, perform the elementwise activation locally, and sum partial results in the contracting projection.
For a standard Transformer layout, this often means roughly two communication collectives per layer: one around the attention output and one around the MLP output. Exact counts depend on the architecture, normalization placement, sequence-parallel choices, and implementation.
Notice what TP does not do: it does not automatically divide the KV cache perfectly. With grouped-query attention, the number of KV heads can be smaller than the number of GPUs, forcing some KV heads to be replicated or requiring a different sharding scheme.
Worked example: why tiny decode batches are awkward¶
Take an illustrative 32-layer Transformer with hidden dimension \(d=4096\), BF16 activations, and four-way tensor parallelism.
Suppose a decode step processes one new token for each of 32 requests. A full hidden-state tensor has
For a ring all-reduce over \(p\) GPUs, the approximate data moved per GPU is
With \(p=4\), that is
If there are two such collectives per layer, across 32 layers, total communication is approximately
24 MiB is not the whole story. Those are 64 separate collectives. Each one has synchronization and startup latency, even when its payload is small. A ring algorithm also requires multiple communication rounds; alternative collective algorithms trade latency against bandwidth differently.
For batch size 1, the hidden-state payload is only 8 KiB. The total communicated bytes become small, but the repeated synchronization can still be expensive. This is one reason tensor parallelism may hurt single-request latency for a model that already fits on one GPU.
During prefill, by contrast, thousands of token positions can be processed together. Communication payloads grow, but so does the matrix-multiplication work available to split across GPUs. The relative economics can be much better.
These calculations are illustrative: real kernels may fuse operations, overlap transfers with computation, use different collectives, or shard activations differently.
Extending yesterday's roofline model¶
With one GPU, we considered compute and high-bandwidth-memory limits. With tensor parallelism, there is a third ceiling: interconnect bandwidth and collective latency.
A rough lower bound becomes
plus any synchronization or communication startup costs that cannot be hidden behind useful work.
That is why the topology matters. GPUs linked by a fast, low-latency interconnect can make fine-grained TP practical. Spreading the same tensor-parallel group across slower network links may turn every layer into a communication bottleneck.
TP versus the other ways to scale¶
Data parallelism: complete model replica per worker
→ more independent requests or training batches
Tensor parallelism: split matrices within each layer
→ GPUs cooperate on the same token/layer
Pipeline parallelism: split layers into consecutive stages
→ activations move between stagesIn training, data-parallel replicas must also coordinate gradients, while tensor-parallel layers participate in both forward and backward communication. In inference, data parallelism is often the simplest way to increase throughput if a complete model replica fits. TP becomes especially valuable when model size, latency goals, or available hardware make one-GPU execution impractical.
These strategies are frequently combined: for example, four GPUs cooperate on each model replica using TP, while several such groups serve independent requests using data parallelism.
Mental model: four accountants and one shared spreadsheet¶
Imagine a large financial model. Four accountants can each calculate a different set of columns simultaneously. That is the column-parallel stage: everyone produces a distinct part of the answer.
But the final balance sheet needs totals that combine all four sets of calculations. Now the accountants must exchange and reconcile their numbers. That is the row-parallel reduction.
Adding accountants reduces individual arithmetic work. It does not eliminate the time spent coordinating. If the spreadsheet is tiny, four people may be slower than one; if it is enormous, the parallel work can outweigh the reconciliation cost.
The deeper lesson¶
Tensor parallelism makes a model larger than one GPU executable and can accelerate heavy computation. But it exchanges local computation and memory pressure for distributed coordination.
The practical question is not merely, “How many GPUs can we use?” It is: At this batch size, context length, model architecture, and interconnect speed, where is the new bottleneck?
That's the distributed version of yesterday's roofline lesson—and a useful foundation for understanding why large-scale LLM serving and training are fundamentally systems-engineering problems.