Tensor Parallelism: Splitting Operations Across GPUs

A fintech infrastructure team is deploying a 70-billion-parameter Hindi-English support assistant to handle chat traffic for a large digital-payments app. In fp16, the weights alone take roughly 140 GB — no single GPU in their fleet, not even an 80 GB A100, can hold the model. So they shard it across GPUs using tensor parallelism, the technique that splits individual matrix operations — not whole layers, not whole batches — across devices so that each GPU holds a slice of every weight matrix and the GPUs cooperate on every forward pass. That much is table stakes, and if you have already studied how Megatron-style column- and row-parallel linear layers work for training a model from scratch, you have the mechanism. What that earlier picture leaves out is what changes when the model is not being trained but served — when thousands of users are mid-conversation, latency is the product, and the model was built with a memory-saving trick called grouped-query attention that interacts with tensor parallelism in a way that quietly caps how much a team can gain by adding more GPUs. That interaction, and the arithmetic behind it, is this chapter.

During a festival-season traffic spike the same team doubles their concurrent chat batch overnight. Their GPU count does not change, but their memory budget and their latency budget both move — and not in the way a naive "more parallelism, more capacity" intuition predicts. Understanding why requires going one level below the weight-matrix view of tensor parallelism, down to how it interacts with attention heads specifically, and with the key/value cache that dominates memory during generation.

A one-paragraph recap: what a tensor-parallel matmul actually does

Shoeybi et al. (Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, NVIDIA, 2019) showed that a single large matrix multiplication Y = XW can be split two ways without any extra communication mid-computation: column-parallel, where W is sliced by columns so each GPU produces an independent slice of Y, and row-parallel, where W is sliced by rows so each GPU produces a partial sum that must be all-reduced (summed across GPUs) before Y is complete. A transformer block is built by chaining a column-parallel projection into a row-parallel one, so the all-reduce happens exactly once per block-half — once after the attention output projection, once after the feed-forward down-projection. If that is new to you, treat it as the substrate; everything below builds on top of it rather than re-deriving it.

Splitting attention by head, and why decode is where it hurts

Multi-head attention is unusually friendly to this scheme. Each head computes its own query, key, and value vectors, its own softmax, and its own weighted sum, entirely independently of every other head — heads only interact when their outputs are concatenated and passed through the output projection. That makes "one GPU, several whole heads" a natural column-parallel split: with tensor-parallel degree P, each GPU is assigned num_heads / P complete heads, runs full local attention for them, and only needs to talk to the other GPUs once, at the row-parallel output projection.

Training-time tensor parallelism and inference-time tensor parallelism run the same split, but under very different arithmetic conditions. During prefill — processing the user's prompt — the matmuls are large (sequence length times hidden size), so they have high arithmetic intensity and the fixed cost of an all-reduce is a small fraction of the work. During decode — generating the reply one token at a time — each step multiplies a comparatively tiny activation (batch size times hidden size, not sequence length times hidden size) against the same huge weight matrices, so the GPUs are memory-bandwidth-bound and the KV cache, not the weights, becomes the dominant memory consumer (Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023). That shift — from compute-bound prefill to memory-bound, KV-cache-dominated decode — is exactly the regime where the interaction below starts to matter.

The GQA ceiling: a worked memory audit

Modern serving-scale models rarely give every attention head its own keys and values. Grouped-query attention (Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, Google Research, EMNLP 2023) has many query heads share a much smaller set of key/value heads, because the KV cache — not the query projection — is what has to be stored per token, per layer, per request, and shrinking it is what lets a server hold more concurrent conversations in GPU memory. Llama-2-70B's published configuration (Touvron et al., Llama 2: Open Foundation and Fine-Tuned Chat Models, Meta AI, 2023) is a clean, real example: 64 query heads, but only 8 key/value heads, with hidden size 8192, feed-forward size 28672, and 80 layers.

Tensor parallelism hands out whole heads to GPUs. At TP = 8, that is a clean split: 64 query heads ÷ 8 = 8 per GPU, and 8 KV heads ÷ 8 = exactly 1 KV head per GPU. Every GPU is self-sufficient. But push the tensor-parallel degree past the KV head count — say to TP = 16, a reasonable move if you also want to halve the weight memory per GPU — and the KV side runs out of heads to distribute. 64 ÷ 16 = 4 query heads per GPU, still fine, but 8 KV heads ÷ 16 GPUs is not an integer. Megatron-style tensor parallelism does not shard a single head's key/value vectors across GPUs, so the only option is replication: each of the 16 GPUs still holds one full KV head, but now two GPUs (the two that each hold half of the original 8-query group sharing that head) hold an identical copy. The system as a whole is now storing 16 copies of what is really only 8 unique heads' worth of data.

The code below computes exactly how this plays out in gigabytes per GPU, using Llama-2-70B's real configuration, fp16 weights, a concurrent batch of 32 requests, and a 4096-token context per request:

# Llama-2-70B-style transformer config (Touvron et al., 2023)
d_model = 8192
d_ff = 28672
n_layers = 80
n_heads = 64
n_kv_heads = 8
head_dim = d_model // n_heads          # 128
bytes_per_param = 2                     # fp16

def params_per_layer():
    ffn = 3 * d_model * d_ff            # SwiGLU: gate, up, down
    kv_dim = n_kv_heads * head_dim
    attn = 2 * d_model * d_model + 2 * d_model * kv_dim   # Q, O, K, V
    return ffn + attn

def weight_gb_per_gpu(tp):
    total_params = params_per_layer() * n_layers
    return total_params * bytes_per_param / tp / 1e9

def kv_cache_gb_per_gpu(tp, batch, seq_len):
    # once tp exceeds n_kv_heads, heads must be replicated: every GPU
    # still ends up storing exactly one full KV head, never a fraction
    heads_per_gpu = 1 if tp > n_kv_heads else n_kv_heads // tp
    bytes_per_token = 2 * heads_per_gpu * head_dim * bytes_per_param  # K and V
    return bytes_per_token * n_layers * batch * seq_len / 1e9

for tp in (8, 16):
    w = weight_gb_per_gpu(tp)
    kv = kv_cache_gb_per_gpu(tp, batch=32, seq_len=4096)
    print(f"TP={tp}: weights={w:.2f} GB/GPU, KV cache={kv:.2f} GB/GPU")

# Output:
# TP=8: weights=17.11 GB/GPU, KV cache=5.37 GB/GPU
# TP=16: weights=8.56 GB/GPU, KV cache=5.37 GB/GPU

Trace it by hand and the output checks out: params_per_layer() sums the feed-forward block (3 × 8192 × 28672 = 704,643,072 params for the gate, up, and down projections) and the attention block (2 × 8192² for Q and O, plus 2 × 8192 × 1024 for K and V, since the KV projection only needs to produce 8 heads × 128 dims = 1024 output features, not 8192) for 855,638,016 params per layer, times 80 layers = 68,451,041,280 params total. At fp16, TP = 8 gives 68.45B × 2 bytes ÷ 8 ÷ 1e9 = 17.11 GB of weights per GPU; doubling to TP = 16 exactly halves it to 8.56 GB, because weight matrices really do split as finely as you like. The KV cache does not. At TP = 8 and TP = 16 alike, heads_per_gpu evaluates to 1 (the branch condition tp > n_kv_heads is false at 8, true at 16, but both routes land on 1 head per GPU), so the per-GPU KV cache is identical: 2 (K and V) × 1 head × 128 dims × 2 bytes = 512 bytes per token per layer, × 80 layers × 32 requests × 4096 tokens = 5.37 GB, unchanged.

Add the two together and the picture sharpens: total memory per GPU falls from 17.11 + 5.37 = 22.48 GB at TP = 8 to 8.56 + 5.37 = 13.93 GB at TP = 16 — a real reduction, about 38%, but far short of the roughly 50%-or-better a naive "doubled the GPUs, should roughly halve the memory" intuition would predict, because more than a third of the TP = 16 total (5.37 of 13.93 GB) is KV cache that TP = 16 did nothing to shrink. And the ceiling does not lift by going further: at TP = 32, heads_per_gpu is still 1, so KV cache per GPU is still 5.37 GB while weights keep falling (4.28 GB) — you can buy weight-memory headroom past TP = 8 almost indefinitely, but KV-cache headroom stops the moment TP passes the KV head count.

Misconception: "more GPUs always means faster generation"

The natural next move, having seen TP = 16 buy real weight-memory savings, is to assume you should push the tensor-parallel degree as high as your GPU budget allows. That is the misconception this section corrects: past a point, adding tensor-parallel ranks makes each decode step slower, not faster, because the all-reduce cost does not shrink anywhere near as fast as the per-GPU compute does.

Use the ring all-reduce cost model: for a message of size M bytes split across P GPUs on a fully connected ring, time ≈ 2(P−1)/P × (M ÷ B), where B is per-link bandwidth. Take the residual-stream activation that has to be summed after the output projection: batch 32 × hidden size 8192 × 2 bytes (fp16) = 524,288 bytes, and NVIDIA's published A100 NVLink bandwidth of 600 GB/s. At TP = 8: 2×7/8 × (524,288 ÷ 600e9) ≈ 1.53 microseconds per all-reduce; there are two all-reduces per layer (one after attention, one after the feed-forward down-projection) across 80 layers, so total communication per decode step ≈ 245 microseconds. At TP = 16, the coefficient 2(P−1)/P barely moves — 1.875 instead of 1.75 — so total communication rises only slightly, to about 262 microseconds.

Now compare that to compute. The linear-layer FLOPs for one decode step across a batch of 32 is roughly 2 × 68.45B params × 32 ≈ 4.38 × 10¹² FLOPs (the standard "2 × params per token" approximation, counting only the matmuls tensor parallelism actually splits). Split across 8 GPUs and run at an assumed achieved throughput of about 90 TFLOP/s per A100 (roughly 58% of its 156 TFLOPS dense fp16 tensor-core peak — a realistic utilization figure for well-tuned inference GEMMs, not the unreachable peak), that is about 6.08 milliseconds per GPU. Communication of 245 microseconds against that is a 4.0% overhead. Halve the per-GPU compute by doubling to TP = 16 (3.04 ms) while communication barely moved (262 μs), and the overhead more than doubles to 8.6%. Keep pushing TP degree and the trend continues: compute per GPU keeps halving while the all-reduce cost is nearly flat, so the fraction of every decode step spent waiting on communication keeps climbing — and that is before accounting for the fact that once TP crosses a single node's boundary, GPUs must talk over InfiniBand or Ethernet at a fraction of NVLink's bandwidth, which can make the crossover far sharper than this same-node estimate shows. Narayanan et al. (Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC21, 2021) recommend keeping the tensor-parallel group inside one NVLink-connected node for exactly this reason, and reserving pipeline or data parallelism for scaling across nodes.

Active recall

Work through these before reading the answers.

1. Why does splitting attention by head require no communication until the very end of the attention block, while a row-parallel feed-forward layer needs an all-reduce partway through?
2. Using the Llama-2-70B configuration above, what is the fp16 weight memory per GPU at TP = 4?
3. At TP = 32, how many query heads and how many KV heads does each GPU hold, and what happens to KV-cache replication compared to TP = 16?
4. The worked memory audit used batch = 32. If a festival-traffic spike pushes concurrent batch to 64, trace every quantity that changes: weight memory per GPU, KV cache per GPU, total memory per GPU at TP = 8 and TP = 16, and the communication-to-compute overhead percentage.
5. True or false: you should always pick the highest tensor-parallel degree your GPU budget allows, since more GPUs means more parallel compute. Justify your answer with a number from this chapter.
6. Why is it specifically the KV cache, and not the query projection, that hits a replication ceiling under tensor parallelism?

Answers

1. Every head's query, key, value, softmax, and weighted sum are computed entirely within that head — no head needs any other head's numbers to produce its own output. Splitting by head is therefore a column-parallel split: each GPU's slice of the result is already complete and correct on its own, and the shards are simply concatenated. The feed-forward down-projection, by contrast, is row-parallel: each GPU multiplies its slice of the weight matrix by the full input and produces only a partial sum toward the final output, so the partial sums from every GPU must be added together (all-reduced) before the result is complete.

2. Total transformer-block parameters are 68,451,041,280. At fp16 (2 bytes) and TP = 4: 68,451,041,280 × 2 ÷ 4 ÷ 1e9 = 34.23 GB per GPU.

3. 64 query heads ÷ 32 GPUs = 2 query heads per GPU — still a clean integer split. But 8 KV heads ÷ 32 GPUs is not an integer, so replication kicks in exactly as it did at TP = 16: each GPU still holds one full KV head (there is no finer unit than "one head" to hand out), so the system stores 32 copies of 8 unique heads — a replication factor of 4×, worse than TP = 16's 2×, for zero additional per-GPU KV-cache saving. Weight memory keeps shrinking regardless (68.45B × 2 ÷ 32 ÷ 1e9 ≈ 4.28 GB/GPU), so the gap between what weights save you and what KV cache refuses to save you keeps widening.

4. Weight memory per GPU is unaffected by batch size — it stays 17.11 GB (TP = 8) and 8.56 GB (TP = 16), since weights don't depend on how many requests are being served. KV cache scales linearly with batch, so it doubles at both TP values: 5.37 → 10.74 GB per GPU, identical for TP = 8 and TP = 16 (the ceiling doesn't move, it just gets bigger in absolute terms). Total memory per GPU becomes 17.11 + 10.74 = 27.85 GB at TP = 8 and 8.56 + 10.74 = 19.30 GB at TP = 16 — the ratio of TP16-total to TP8-total rises from about 0.62 (at batch 32) to about 0.69 (at batch 64), meaning the relative benefit of the higher TP degree actually shrinks as traffic grows, because the KV-cache floor is claiming a larger share of the total. The communication-to-compute overhead percentage, however, does not change with batch: both the all-reduce message size and the FLOPs scale linearly with batch, so the ratio between them cancels out — it stays at roughly 4.0% (TP = 8) and 8.6% (TP = 16) at batch 64 exactly as at batch 32. Batch size reshapes the memory economics of choosing a TP degree, but leaves the latency-overhead economics untouched.

5. False. The comm/compute overhead derivation above shows it roughly doubling, from about 4.0% to about 8.6%, just going from TP = 8 to TP = 16 within a single NVLink-connected node — because per-GPU compute keeps halving while the all-reduce cost barely drops. Cross a node boundary and the effect compounds further, since inter-node bandwidth is typically far below NVLink's 600 GB/s. Past some degree, adding more tensor-parallel ranks makes each decode step slower rather than faster, which is why production systems cap TP at a node's GPU count and use other parallelism strategies to scale beyond it.

6. Grouped-query attention deliberately gives far fewer heads to the keys and values than to the queries, specifically to shrink the KV cache — the piece of memory that has to be stored per token, per request, and that dominates memory during decode. Tensor parallelism, however, only knows how to distribute whole heads to GPUs; it has no mechanism (in the standard Megatron scheme) for splitting a single head's key/value vectors further. Query heads have 64 of themselves to spread around, so they don't hit this wall until TP = 64. KV heads have only 8, so the wall arrives at TP = 9 and every degree beyond it pays for replication instead of getting a real memory dividend.

Think About It

Think about this: How would you explain tensor parallelism: splitting operations across gpus to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where tensor parallelism: splitting operations across gpus is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting tensor parallelism: splitting operations across gpus to at least 3 other topics you have studied.