Personalize this lesson
Adapt explanations and teaching visuals to your background and preferred voice.
A large language model (LLM) has an estimated 70 GB weight payload. Give each GPU an assumed 80 decimal GB budget and set aside 16 GB for other allocations and headroom. The weights exceed the remaining 64 GB before any request arrives. A runtime may reject this configuration or run out of memory. A second GPU adds capacity, but which tensors belong on each device, and what must cross the wires to recover the model's mathematical result?
Earlier chapters treated a single GPU as a standalone serving box: weights load into VRAM, and then active requests compete for whatever bytes remain for their KV cache. Data-parallel replicas help when a full model copy already fits on one device. They can't shrink an oversized model that doesn't fit in the first place.
Model parallelism divides one model across accelerators. Communication, replicated tensors, stage imbalance, and idle slots affect its benefit. In a coding assistant, weight placement helps a replica load; per-rank KV state and execution work affect which sessions can be admitted within latency targets. Matching the unsharded mathematical computation doesn't guarantee bit-identical floating-point results or sampled text.
Use Qwen3.6-35B-A3B's approximately 35B total and 3B active parameters as the starting estimate.[1] At two BF16 bytes per parameter, that gives about 70 decimal GB. MoE routing selects only some experts per token, but a GPU-resident deployment still stores the full expert set across its devices. Offloading changes that placement and adds transfers. Active parameters are a compute proxy, not a full FLOP count or weight-memory estimate.
We'll use 70 GB as our rounded weight checkpoint, with an assumed 16 GB runtime reserve per device. Real production setups replicate some layers, balance shards unevenly, and load multimodal encoders. The matrix examples that follow teach the core sharding mechanics from first principles.
Why can Qwen3.6-35B-A3B need more than one GPU before serving any real traffic?
Answer
Under the fixture's assumed 80 GB budget and 16 GB reserve, only 64 GB remains for an approximately 70 GB BF16 payload. Actual sizing must include the checkpoint, KV pool, workspace, replicated tensors, and headroom on every rank. The 3B active count doesn't make the full weight payload 6 GB. A different precision or offload path needs another budget.

Why inference sharding differs from training
Distributed training focuses on gradients, optimizer states, backward passes, and aggregate throughput across thousands of sample steps. Serving runs on a very different clock: time to first token (TTFT) exposes prompt processing latency, tokens per second (TPS) governs the autoregressive decode loop, and per-token KV cache memory controls concurrency.
The names of parallel techniques carry over from training literature, but their operational costs change completely. An optimization that saves activation memory during a backward pass might do nothing for inference latency, while an engine layout that fits a single prompt might leave pipeline hardware starved for work. Here's how their trade-offs separate:
| Technique | Training concern | Inference concern |
|---|---|---|
| Tensor parallelism (TP) | Split matmuls and gradients | Split weights and activations, paying for intra-layer all-reduce collectives |
| Pipeline parallelism (PP) | Fill stages with microbatches | Balance layer stages while avoiding decode latency bloat |
| Sequence parallelism | Shard LayerNorm and activations alongside TP | Training memory optimization, distinct from inference long-context serving |
| Context parallelism | Split long training documents across ranks | Distribute massive prompt attention and KV caches across devices |
| Data parallelism | Replicate model state and synchronize gradients | Independent serving replicas; EP deployments can couple their expert work |
Serving engines like vLLM expose tensor_parallel_size and pipeline_parallel_size. Checked September 22, 2026: its guide suggests one GPU when sufficient, intra-node TP for a larger model, and TP plus PP across nodes. It also describes intra-node PP for uneven splits or GPUs without NVLink, and allows TP across nodes.[2] These are candidates to measure, not transport requirements or speed guarantees.
A logical plan doesn't specify physical transport. TP=4 assigns work to four ranks; it doesn't reveal their topology, shared paths, achieved collective latency, or whether communication is exposed on the request's critical path.
The NVIDIA Collective Communication Library (NCCL) detects hardware topology and dispatches optimized collective primitives like all-reduce, all-gather, and all-to-all across available paths.[3] Always track the physical interconnect alongside your logical topology: a TP=4 layout over NVLink behaves very differently from TP=4 over PCIe.
Before picking a strategy, write down your per-device memory ledger:
1weight memory ~= parameters x bytes per parameter
2serving memory ~= weights + KV cache + runtime buffers + safety headroomUse two bytes per BF16 parameter for raw payload. A preallocated KV pool may keep its device allocation constant while occupied blocks change. The following even-shard calculation gives a lower bound under its assumed reserve, before KV, replication, and placement constraints:
1from math import ceil
2
3def weight_memory_gb(params_billion: float, bytes_per_param: float) -> float:
4 return params_billion * bytes_per_param
5
6params_billion = 35
7gpu_gb = 80
8reserve_fraction = 0.20
9usable_gb = gpu_gb * (1 - reserve_fraction)
10
11bf16_weights = weight_memory_gb(params_billion, bytes_per_param=2)
12raw_min_gpus = ceil(bf16_weights / usable_gb)
13int4_weights = weight_memory_gb(params_billion, bytes_per_param=0.5)
14
15print(f"Qwen3.6-35B-A3B BF16 total weights: {bf16_weights:.0f} GB")
16print(f"80GB GPU usable with 20% reserve: {usable_gb:.0f} GB")
17print(f"minimum GPUs for weights with reserve: {raw_min_gpus}")
18print(f"Qwen3.6-35B-A3B INT4 total weight-only estimate: {int4_weights:.1f} GB")1Qwen3.6-35B-A3B BF16 total weights: 70 GB
280GB GPU usable with 20% reserve: 64 GB
3minimum GPUs for weights with reserve: 2
4Qwen3.6-35B-A3B INT4 total weight-only estimate: 17.5 GBThe two-device bound belongs to this budget, not every deployment. Legal sharding depends on the runtime and architecture. INT4's 17.5 GB raw estimate excludes scales, packing, and unquantized tensors. The official Qwen card gives a TP=8 launch example at 262,144 context; it doesn't prove eight GPUs are required or identify KV as the only reason for that choice.[1]
Calculate the same fixture's per-rank remainder:
1weights_gb = 70
2gpu_gb = 80
3reserved_runtime_gb = 16
4
5for tp_size in (1, 2):
6 shard_gb = weights_gb / tp_size
7 kv_room_gb = gpu_gb - reserved_runtime_gb - shard_gb
8 print(f"TP={tp_size}: {shard_gb:.0f} GB weights/GPU, {kv_room_gb:.0f} GB left for KV")1TP=1: 70 GB weights/GPU, -6 GB left for KV
2TP=2: 35 GB weights/GPU, 29 GB left for KVTensor parallelism
Tensor parallelism shards operations inside a layer, whether or not that individual layer is too large. Each rank computes with local weight slices and communicates the required results. The participating ranks jointly advance the batch; that doesn't mean all devices perform useful work at every instant.
A rank is an individual execution worker, typically one GPU process. A collective coordinates communication across all ranks in the group. The Megatron-LM formulation makes this partitioning elegant by pairing complementary matrix cuts.[4]
Consider a standard linear projection:
1y = xWIf we slice across ranks, each chip computes a partial result. Communication synchronizes those fragments so downstream layers receive the correct mathematical output. Tensor parallelism pays off when GPUs connect over ultra-fast links like NVLink.
The Megatron pattern: column then row
The original Megatron-LM pairing avoids an intermediate communication by matching column and row cuts.[4] It reduces communication in the block; it doesn't guarantee communication is no longer a bottleneck.
In the MLP block :
- Column-parallel expansion: Split weight matrix along its columns into . The input vector is replicated on all ranks. Rank 0 calculates while Rank 1 calculates . Because GeLU is an element-wise function, each rank applies to its local activation slice independently. No cross-device communication happens here. If you had split by rows, you'd need an immediate all-reduce before the activation, because .
- Row-parallel projection: Split weight matrix along its rows into . The local activation from Rank 0 () has the exact inner dimension required to multiply directly against local row shard . Rank 0 produces partial product , and Rank 1 produces .
- Single all-reduce reduction: Notice that . Slicing by rows turns the final output into a sum of partial products. One sum all-reduce collective combines these partial outputs across ranks, and the bias is added once to restore the exact logical output.
The multi-head attention block mirrors this exact column-then-row symmetry. The query (), key (), and value () projections shard column-wise, which maps cleanly onto independent attention heads. Each rank computes scaled dot-product attention for its local subset of heads without talking to its neighbors. The attention output projection () then shards row-wise, taking local head outputs and producing partial products. A single sum all-reduce restores the full multi-head output.
For this classic dense forward pattern, count two sum all-reduces per transformer layer. This excludes embeddings/output handling and assumes TP>1. Sequence/context parallelism, expert routing, fusion, and other runtime layouts can change the collectives; it isn't a universal count for the hybrid Qwen model.
![The two-rank MLP fixture replicates input [1,2], column-shards A, applies local GeLU, and row-shards B. Summing partials and adding bias once gives approximately [-2.6585,8.3411]. The output shape is [1,2]; its values differ from the input.](/cdn/content-image/fundamentals/model-parallelism-for-llm-inference/illustrations/_generated/tensor_shards_dark.png?v=85796629cc3c)
Verify the split with actual tensors
Let's test this sharding math directly in PyTorch. We'll take an input with two features, expand it to four hidden features through matrix , apply GeLU, and project back to two features through matrix . GPU 0 owns the first two hidden features; GPU 1 owns the last two. Their local outputs must be summed, not averaged or concatenated:
1import torch
2import torch.nn.functional as F
3
4torch.set_num_threads(1)
5x = torch.tensor([[1., 2.]], dtype=torch.float64)
6A = torch.tensor([[1., -1., 2., 0.], [0., 1., -1., 2.]], dtype=x.dtype)
7B = torch.tensor([[1., 0.], [0., 1.], [1., 1.], [-1., 2.]], dtype=x.dtype)
8bias = torch.tensor([0.5, -0.5], dtype=x.dtype)
9dense = F.gelu(x @ A) @ B + bias
10
11def sharded_mlp(x, A, B, bias, ranks):
12 if type(ranks) is not int or ranks < 1 or ranks > A.shape[1] or A.shape[1] % ranks:
13 raise ValueError("hidden features must divide evenly across ranks")
14 if A.shape[1] != B.shape[0]:
15 raise ValueError("hidden dimensions disagree")
16 partials = [F.gelu(x @ a) @ b for a, b in
17 zip(A.chunk(ranks, dim=1), B.chunk(ranks, dim=0))]
18 return torch.stack(partials).sum(dim=0) + bias, partials
19
20sharded, partials = sharded_mlp(x, A, B, bias, ranks=2)
21torch.testing.assert_close(sharded, dense)
22wrong_mean = torch.stack(partials).mean(dim=0) + bias
23wrong_bias = sum(part + bias for part in partials)
24assert not torch.allclose(wrong_mean, dense)
25assert not torch.allclose(wrong_bias, dense)
26print("xA:", (x @ A).tolist())
27print("dense output:", [round(v, 4) for v in dense[0].tolist()])
28print("sharded equals dense:", torch.allclose(sharded, dense))
29print("mean reduction and duplicated bias: rejected")1xA: [[1.0, 1.0, 0.0, 4.0]]
2dense output: [-2.6585, 8.3411]
3sharded equals dense: True
4mean reduction and duplicated bias: rejectedWith two ranks, each local weight slice has shape [2, 2], each hidden activation slice is [1, 2], and each local projection slice has shape [2, 2]. Both ranks compute their local contribution independently. An all-reduce sum combines the slices so every participating rank holds the identical final activation vector.[3]
KV head counts and sharding constraints
For clean head sharding, your tensor parallel size must divide the attention head counts. In standard multi-head attention (MHA), that means .
Under Grouped-Query Attention (GQA), multiple query heads share a single key-value head. Slicing KV heads across ranks without replicating them imposes a much stricter condition:
For example, Llama 3 70B has 64 query heads and 8 KV heads. On TP=8, each GPU owns 8 query heads and exactly 1 KV head (). On TP=4, each GPU owns 16 query heads and 2 KV heads (). Both divide evenly.
What happens when TP exceeds the KV-head count? Qwen's ten full-attention layers have 16 query heads, 2 KV heads, and head dimension 256. Its other thirty layers use Gated DeltaNet state; a full-history KV formula doesn't describe all forty layers.[5]
1q_heads = 16
2kv_heads = 2
3
4for tp in (1, 2, 4, 8):
5 divides_q = q_heads % tp == 0
6 divides_kv = kv_heads % tp == 0
7 print(f"TP={tp}: divides Q heads={divides_q}, divides KV heads={divides_kv}")1TP=1: divides Q heads=True, divides KV heads=True
2TP=2: divides Q heads=True, divides KV heads=True
3TP=4: divides Q heads=True, divides KV heads=False
4TP=8: divides Q heads=True, divides KV heads=FalsePure nonduplicating head sharding fails that divisibility test at TP=4 or 8. Current vLLM's QKV linear layer instead supports KV replication when TP is a multiple of the smaller KV-head count.[6] At TP=8 with two KV heads, four ranks carry each head's projection. Verify the model/backend's cache layout too.
For plain replicated-head full-history caching, TP=2 to TP=8 still leaves one KV head per rank. That history term stops shrinking; the freed weight bytes can nevertheless enlarge the available pool. Other hybrid-layer state has separate rules. Decode context parallelism can shard history within those replicated groups when supported, so zero further KV relief is not a universal limit.
Count communication, then measure its exposed cost
Under the classic 80-layer dense pattern above, one batch decode iteration has 160 in-layer all-reduces. Embedding and sampling collectives can add to the count.
Across 128 decode iterations, that assumed pattern exceeds 20,000 in-layer all-reduces. A 128-token completion normally needs only 127 decode iterations after prefill produces its first token:
1layers = 80
2all_reduces_per_layer = 2
3decode_iterations = 128
4
5per_step = layers * all_reduces_per_layer
6generation_total = per_step * decode_iterations
7
8print(f"all-reduces per batch decode iteration: {per_step}")
9print(f"all-reduces over {decode_iterations} iterations: {generation_total:,}")
10print(f"decode-only count for 128 output tokens: {per_step * 127:,}")1all-reduces per batch decode iteration: 160
2all-reduces over 128 iterations: 20,480
3decode-only count for 128 output tokens: 20,320Dependencies between blocks can expose communication on the critical path. A sum of collective durations is not automatically the exposed delay: overlap, fusion, algorithms, and launch behavior matter. First examine a deliberately serial startup model:
Compare two assumed per-call startup costs:
1collectives_per_token = 160
2
3for assumed_collective_latency_us in (5, 20):
4 floor_ms = collectives_per_token * assumed_collective_latency_us / 1000
5 print(f"{assumed_collective_latency_us} us startup -> {floor_ms:.1f} ms/step before payload transfer")15 us startup -> 0.8 ms/step before payload transfer
220 us startup -> 3.2 ms/step before payload transferThe 5 and 20 microseconds are authored assumptions, not NVLink/PCIe/network measurements. In a fully serial comparison, adding 3.2 ms to a 6 ms baseline raises duration by 53.3%. It doesn't prove pure idle time or a particular fabric's latency. Measure the operation at the relevant rank count, payload, dtype, and topology; NCCL tests distinguish small-message latency, algorithm bandwidth, and corrected bus bandwidth.[7]
Now check payload transfer time for a batch of 8 requests in BF16:
1batch, hidden, dtype_bytes = 8, 4096, 2 # BF16 activations
2tp = 4
3collectives_per_token = 160
4# Ring all-reduce bytes ≈ 2*(TP-1)/TP * activation_bytes
5activation_bytes = batch * hidden * dtype_bytes
6ring_factor = 2 * (tp - 1) / tp
7bytes_per_collective = ring_factor * activation_bytes
8bytes_per_token = bytes_per_collective * collectives_per_token
9# Assumed sustained per-rank send bandwidth, not aggregate send + receive.
10link_gbs = 300
11payload_ms = bytes_per_token / (link_gbs * 1_000_000_000) * 1000
12
13print(f"bytes sent per rank per collective: {bytes_per_collective / 1024:.1f} KiB")
14print(f"payload-only estimate @ {link_gbs} GB/s: {payload_ms:.2f} ms/step")
15print("Compare to startup-only estimates (0.8 to 3.2 ms); larger batch raises payload.")1bytes sent per rank per collective: 96.0 KiB
2payload-only estimate @ 300 GB/s: 0.05 ms/step
3Compare to startup-only estimates (0.8 to 3.2 ms); larger batch raises payload.The 0.05 ms value uses an assumed 300 GB/s sustained per-rank send rate and a ring traffic model. It isn't an observed NVLink duration. Comparing it with the separate startup assumptions suggests what to profile; summing the two models still doesn't predict total exposed latency. NCCL may choose a tree or another algorithm.
TP can cross nodes, and a fast scale-up domain needn't stop at eight GPUs: NVIDIA GB200 NVL72 connects 72 GPUs by NVLink.[8] Domain size doesn't guarantee that a particular TP degree divides a model or improves its performance.
Roughly how much communication does tensor parallelism add per decode step, and why does the interconnect decide if it helps?
Answer
The classic dense pattern counts 160 in-layer all-reduces for an 80-layer iteration, independently of batch size. Assuming 5 or 20 us startup per call gives 0.8 or 3.2 ms in a serial startup-only model. Those values don't characterize a real interconnect. Measure exposed communication against compute savings, including payload, overlap, and any extra collectives.
Pipeline parallelism
Pipeline parallelism partitions the model along its depth. GPU 0 hosts layers 0 to 39, while GPU 1 hosts layers 40 to 79. Activations travel forward across stage boundaries: GPU 0 computes its assigned layers and forwards the intermediate hidden state to GPU 1.
Classic PP sends hidden activation tensors at stage boundaries, plus any required metadata. This reduces communication frequency compared with the earlier dense TP pattern. TP reductions also communicate activations, rather than whole weight matrices. Fewer boundary transfers can make PP useful over slower links; measure their size, scheduling, and imbalance costs.
For ordinary autoregressive decode, the next token depends on the previous token's final logits and sampling. One request must cross all stages. Other independent microbatches can occupy earlier stages; a shallow queue leaves those slots idle. Speculative/multi-token methods need a separate schedule.
Why can pipeline parallelism hurt single-request latency?
Answer
A request crosses every stage, with transfers and scheduling delays between them. Its compute time is the sum of the stage times. Dividing the original layer stack into p pieces doesn't multiply the original unsharded latency by p. Idle stage slots limit utilization; transfers, imbalance, and changed kernels determine the latency difference.
Microbatch scheduling and pipeline bubbles
When multiple independent requests arrive, the serving engine groups them into microbatches and interleaves their execution across stages. Stage processes microbatch during time slot .
Assume equal-duration stages, independent equal-size microbatches, one forward traversal, and no transfer/queue cost. Then stages and microbatches take slots. Occupied slots total out of . This is a finite fill/drain fixture, not a complete steady-state autoregressive serving schedule.
The pipeline bubble fraction is the proportion of idle hardware time during the fill (warm-up) and drain (cool-down) phases:
Relative to the ideal throughput interval , fill/drain overhead has this dimensionless ratio:
The extra elapsed time is . Neither ratio is a multiplier on the original full-model single-request latency.

Let's compute the occupancy schedule explicitly:
1def ideal_pipeline_utilization(stages: int, microbatches: int) -> float:
2 if (type(stages) is not int or type(microbatches) is not int
3 or stages < 1 or microbatches < 1):
4 raise ValueError("stages and microbatches must be positive integers")
5 return microbatches / (microbatches + stages - 1)
6
7def schedule(stages, microbatches):
8 ideal_pipeline_utilization(stages, microbatches)
9 return [[m if 0 <= (m := t - s) < microbatches else None
10 for t in range(microbatches + stages - 1)]
11 for s in range(stages)]
12
13for microbatches in (1, 4, 12):
14 grid = schedule(3, microbatches)
15 occupied = sum(cell is not None for row in grid for cell in row)
16 utilization = ideal_pipeline_utilization(stages=3, microbatches=microbatches)
17 assert occupied == 3 * microbatches
18 assert abs(occupied / sum(map(len, grid)) - utilization) < 1e-12
19 print(f"3 stages, {microbatches:2d} microbatches: {utilization:.1%} ideal utilization")13 stages, 1 microbatches: 33.3% ideal utilization
23 stages, 4 microbatches: 66.7% ideal utilization
33 stages, 12 microbatches: 85.7% ideal utilizationCounts must be discrete. A fractional stage count must not quietly produce a plausible utilization:
1for stages, microbatches in ((2.5, 4), (3, 1.5), (True, 4), (3, True)):
2 try:
3 ideal_pipeline_utilization(stages, microbatches)
4 except ValueError:
5 pass
6 else:
7 raise AssertionError("noninteger pipeline count accepted")
8print("fractional and Boolean stage/microbatch counts rejected")1fractional and Boolean stage/microbatch counts rejectedWhen pipeline parallelism is justified
For the three-stage, one-microbatch fixture, two-thirds of stage slots are idle. Compare its critical path with a full-stack baseline before turning that occupancy into a latency claim:
1stage_compute_ms = (2.0, 2.0, 2.0)
2assumed_transfer_ms = 0.2
3full_stack_ms = sum(stage_compute_ms)
4pipeline_ms = full_stack_ms + (len(stage_compute_ms) - 1) * assumed_transfer_ms
5print(f"assumed full-stack compute: {full_stack_ms:.1f} ms")
6print(f"same compute plus two transfers: {pipeline_ms:.1f} ms")
7print(f"latency ratio: {pipeline_ms / full_stack_ms:.2f}x, not 3x")
8assert abs(pipeline_ms - 6.4) < 1e-121assumed full-stack compute: 6.0 ms
2same compute plus two transfers: 6.4 ms
3latency ratio: 1.07x, not 3xThese are chosen times, not a guarantee that partitioned kernels preserve the baseline compute time.
Pipeline parallelism in inference is justified primarily when:
- Depth partitioning solves a placement constraint. Use PP when a validated layout needs more capacity, uneven layer splits, or fewer frequent collectives. It can be useful inside a node without NVLink, as vLLM's guide notes.[2]
- The topology favors a hybrid. TP groups can stay within a fast domain while PP crosses slower links. TP=8 per node is one candidate, not a universal optimum or domain limit.
- The workload and scheduler sustain useful overlap. More independent microbatches can amortize bubbles; no fixed m≥8 threshold keeps arbitrary stages continuously busy.
The names Llama 3.1 405B and DeepSeek-V3 671B don't by themselves prove a multi-node requirement.[9][10] Check the raw INT4 arithmetic:
1assumed_group_budget_gb = 8 * 80
2for parameters_billion in (405, 671):
3 packed_gb = parameters_billion * 0.5
4 print(f"{parameters_billion}B raw INT4: {packed_gb:.1f} GB vs {assumed_group_budget_gb} GB group budget")
5 assert packed_gb < assumed_group_budget_gb
6print("Raw weights fit this sum; scales, per-rank placement, KV, and latency still need validation.")1405B raw INT4: 202.5 GB vs 640 GB group budget
2671B raw INT4: 335.5 GB vs 640 GB group budget
3Raw weights fit this sum; scales, per-rank placement, KV, and latency still need validation.
Context parallelism for long-sequence serving
Long inputs can make attention work or retained history a constraint alongside weights. Required storage depends on the architecture; sequence sharding doesn't extend a checkpoint's trained/supported context.
In Megatron Core, sequence parallelism shards the sequence dimension across LayerNorm and Dropout operations alongside tensor parallelism to save activation memory during training.[11] It doesn't distribute prompt attention or KV cache state across independent serving nodes.
Context parallelism (CP) partitions the input sequence across multiple GPUs, allowing chips to share the KV cache and attention computation of a single massive prompt.
For a hypothetical 1M-token full-attention GQA history (80 layers, 8 KV heads, head dimension 128), count raw FP16 payload. This isn't a supported Llama context or Qwen's hybrid state:
1tokens = 1_000_000
2layers = 80
3kv_heads = 8
4head_dim = 128
5dtype_bytes = 2
6devices = 4
7
8kv_gib = 2 * tokens * layers * kv_heads * head_dim * dtype_bytes / 1024**3
9print(f"dense GQA 1M-token KV footprint: {kv_gib:.1f} GiB")
10print(f"even {devices}-way context shard: {kv_gib / devices:.1f} GiB/device before overhead")1dense GQA 1M-token KV footprint: 305.2 GiB
2even 4-way context shard: 76.3 GiB/device before overheadAt 305.2 GiB, history alone exceeds one small-device budget. Four even shards still need 76.3 GiB each, before weights/workspace. The previous 80,000,000,000-byte budget is only 74.5 GiB, so four shards don't make that fixture fit. Inspect both prefill and decode mechanisms.
Ring Attention
Ring Attention shards the sequence dimension into blocks across a logical ring of GPUs.[12] Each GPU holds its local Query block throughout the layer.
During each step of the ring:
- GPU computes blockwise self-attention using its local Query block and its current Key-Value block.
- Simultaneously, GPU sends its current Key-Value block to GPU and receives a new Key-Value block from GPU over peer-to-peer links.
- Online softmax merges statistics and weighted outputs across visited blocks. For the simple noncausal scheme, each query block visits C KV blocks, requiring C−1 transfers after its local block; causal masks can skip work.
Transfers can overlap block attention. Fully hiding them requires enough compute and suitable buffering/topology; the paper's overlap condition is not a guarantee for short decode queries.[12]
DeepSpeed Ulysses
DeepSpeed Ulysses was introduced for long-sequence training using all-to-all collectives.[13] Its layout mechanism helps explain prefill attention, but isn't itself a promise that a serving runtime retains a sequence-sharded decode cache.
Starting with sequence-partitioned inputs:
- Each rank first projects its local tokens into Q/K/V. An all-to-all then redistributes those projected tensors from sequence slices (all heads, S/C tokens) to head slices (assigned heads, full sequence).
- Each GPU computes standard local FlashAttention across the full sequence for its assigned subset of heads.
- A second All-to-All collective transposes the output back to sequence-parallel before the linear output projection.
The simple nonduplicating layout needs compatible head divisibility, including the smaller KV-head count under GQA. Supported replication or hybrid layouts change that limit. Ring sequence partitioning doesn't require splitting heads by C, but brings ring transfers, buffer requirements, and masking/load-balance concerns.
Decode context sharding addresses another axis
Current vLLM documentation separates prefill CP from decode context parallelism (DCP).[14] DCP shards retained history inside existing TP groups to reduce KV duplication, without adding another factor to the GPU count. The guide describes GQA/MLA support; validate the particular model and attention backend, especially hybrid state.
Try the ideal layout for two KV heads, TP=8, DCP=4, and 32 history positions. Ignore page rounding and metadata:
1tp, kv_heads, dcp, history = 8, 2, 4, 32
2owned = []
3for rank in range(tp):
4 head = rank // dcp
5 lane = rank % dcp
6 owned.append({(head, token) for token in range(lane, history, dcp)})
7assert set.union(*owned) == {(head, token) for head in range(kv_heads) for token in range(history)}
8assert sum(map(len, owned)) == kv_heads * history
9print(f"plain replicated-head cache: {history} head-position pairs/rank")
10print(f"ideal DCP={dcp}: {len(owned[0])} pairs/rank")
11print(f"total physical pairs: {tp * history} -> {sum(map(len, owned))}")1plain replicated-head cache: 32 head-position pairs/rank
2ideal DCP=4: 8 pairs/rank
3total physical pairs: 256 -> 64These ranks must combine attention statistics and weighted contributions. Averaging independently normalized shard outputs is generally wrong. The CPU exercise checks position ownership, not attention numerics, backend support, or a measured speedup.
When does context-aware parallelism become relevant for inference?
Answer
When a supported long-context workload is limited by attention work or retained history. Prefill query partitioning, Ulysses-style head redistribution, Ring Attention, and decode history sharding have different communication/storage layouts. Check the runtime's specific phase and cache support; a training sequence-parallel label alone doesn't establish it.
Expert parallelism for MoE models
Mixture-of-experts architectures present another dimension of model partitioning: where does each expert live? Instead of routing every token through identical dense feed-forward blocks, a router directs each token to top- experts.[15]
Expert parallelism (EP) places different experts on different GPUs, scaling total parameter capacity without replicating every expert onto every chip.
Tensor parallelism and expert parallelism slice parameters along different axes:
- TP slices expert matrices within a tensor-parallel group.
- EP places different experts on different ranks. A runtime can combine expert placement with tensor sharding or keep redundant expert copies.
Routing requires dispatching token activations to expert owners and combining weighted results back at their origins. This is logically an all-to-all exchange; the transport may use specialized dispatch/combine kernels or all-gather/reduce-scatter rather than two literal NCCL all-to-all calls.[16]
Router decisions depend on hidden representations, not an assumed keyword rule. Skewed assignments can create stragglers; measure expert work and transfers rather than infer latency from counts alone. Checked September 22, 2026: vLLM's --enable-expert-parallel forms an expert group across TP×DP ranks, while attention uses TP inside each DP group. Its load-balancing options can add redundant expert copies.[16]
Here's an explicit simulation showing how token assignments distribute across two devices:
1routes = [(0, 1), (0, 1), (0, 1), (0, 2), (1, 3), (2, 3)]
2inputs = [1., 2., 3., 4., 5., 6.]
3gates = (0.75, 0.25)
4owners = {0: 0, 1: 0, 2: 1, 3: 1}
5queues = [[], []]
6for token, selected in enumerate(routes):
7 for expert, gate in zip(selected, gates):
8 queues[owners[expert]].append((token, expert, gate))
9
10def expert_value(expert, value):
11 return (expert + 1) * value
12
13combined = [0.] * len(inputs)
14for queue in queues:
15 for token, expert, gate in queue:
16 combined[token] += gate * expert_value(expert, inputs[token])
17reference = [sum(g * expert_value(e, x) for e, g in zip(r, gates))
18 for x, r in zip(inputs, routes)]
19assert combined == reference
20loads = list(map(len, queues))
21assert sum(loads) == 2 * len(inputs)
22print("assignments per device:", loads)
23print("combined output:", combined)
24print(f"peak / average assignments: {max(loads) / (sum(loads) / 2):.2f}x")1assignments per device: [8, 4]
2combined output: [1.25, 2.5, 3.75, 6.0, 12.5, 19.5]
3peak / average assignments: 1.33xMulti-GPU serving Pareto trade-offs
Let's synthesize these parallel dimensions into an engineering decision framework. On a cluster of 8x 80 GB GPUs, how should you partition your hardware to serve a 70 GB model under production traffic?
Three distinct operational topologies compete:
1. Pure Tensor Parallelism (TP=8)
- Profile: A single replica spanning all 8 GPUs.
- Potential benefit: More ranks share eligible matrix work and weight storage; measure whether that reduces prompt or decode duration.
- Costs: Collectives, smaller local kernels, replicated state, and a shared replica's admission policy. Neither fastest latency nor maximum communication cost follows from TP=8 alone; the classic 160 count belongs to the earlier dense example.
2. Sharded Data-Parallel Replicas (TP=2, DP=4)
- Profile: Four independent 2-GPU replicas, each running
TP=2. - Potential benefit: Four separately resident pairs can process independent requests if each validated layout fits. Fourfold scaling relative to one such pair requires enough demand, good balance, and no shared bottleneck; it isn't fourfold versus one eight-GPU TP replica.
- Costs: Replicated weight storage, separate caches, load balancing, and skew. Single-request latency can be higher or lower than TP=8. These pairs are independent only if their expert layers are also independent; enabling cross-group EP couples them.
3. Prefill-Decode Disaggregation
- Profile: Physically separated worker pools for the prefill and decode phases.[17][18]
- Potential benefit: Tune each phase's placement and scheduling separately and reduce mixed-phase interference. Suitable parallelism depends on the measured model, prompt/history mix, backend, and hardware.
- Costs: Both queues, KV-state movement/re-layout, compatibility, coordination, and additional workers. Transfers may overlap some work but still need critical-path accounting. No TP degree or RDMA label guarantees the benefit.

Let's verify the KV cache headroom for a TP=2 deployment assuming each concurrent request requires 0.75 GB of KV cache across the pair:
1gpu_count = 2
2gpu_capacity_gb = 80
3weights_gb = 70
4runtime_reserve_gb = 32
5assumed_kv_per_request_gb = 0.75 # total across both GPUs, not per GPU
6
7kv_budget_gb = gpu_count * gpu_capacity_gb - weights_gb - runtime_reserve_gb
8arithmetic_batch_ceiling = int(kv_budget_gb / assumed_kv_per_request_gb)
9
10print(f"KV budget after weights and reserve: {kv_budget_gb} GB")
11print(f"arithmetic request ceiling at assumed KV/request: {arithmetic_batch_ceiling}")
12print("Latency and burst headroom determine the admitted batch below this ceiling.")1KV budget after weights and reserve: 58 GB
2arithmetic request ceiling at assumed KV/request: 77
3Latency and burst headroom determine the admitted batch below this ceiling.What to measure
Compare configurations on the same hardware accounting, model/precision, arrival mix, prompt/output distribution, quality checks, and measurement boundary. Record per-rank allocations and achieved communication behavior alongside user-visible delivery.
Structure your benchmarking decisions through a three-gate evaluation process:
- Gate 1: Physical Fit: Does the checkpoint, runtime reserve, and required KV cache headroom fit in physical VRAM without triggering out-of-memory crashes?
- Gate 2: Latency SLOs: Does the candidate satisfy both p95 TTFT (e.g. ms) and p95 TPOT (e.g. ms) under target concurrency?
- Gate 3: Cost: Among eligible candidates, compare sustained useful output per billed interval. Raw TPS alone doesn't establish request goodput or cost efficiency when GPU counts differ.
First exercise fit and latency filtering with assigned values. This sketch ranks remaining candidates by raw TPS; it has no prices or deployment action:
1benchmarks = [ # synthetic fixture, not measurements
2 {"name": "TP=1", "fits": False, "ttft_p95": None, "tpot_p95": None, "tps": None},
3 {"name": "TP=2", "fits": True, "ttft_p95": 310, "tpot_p95": 45, "tps": 620},
4 {"name": "TP=4", "fits": True, "ttft_p95": 430, "tpot_p95": 62, "tps": 700},
5]
6ttft_limit, tpot_limit = 400, 55
7eligible = [b for b in benchmarks if b["fits"] and b["ttft_p95"] <= ttft_limit and b["tpot_p95"] <= tpot_limit]
8best = max(eligible, key=lambda b: b["tps"], default=None)
9print(f"eligible configurations: {[b['name'] for b in eligible]}")
10print(f"highest-throughput configuration inside SLO: {best['name'] if best else 'none'}")1eligible configurations: ['TP=2']
2highest-throughput configuration inside SLO: TP=2In this synthetic fixture, TP=4's assigned 700 TPS doesn't override its failing latencies. TP=2 is the only eligible candidate under these gates. No trace identifies the cause, and eligibility here isn't production validation. A p95 of per-request average TPOT also doesn't bound each token gap or prove that 95% of requests meet both limits jointly.
Now suppose both candidates pass the other gates. Predict which wins on GPU-only cost when the larger group costs twice as much per hour:
1eligible = [ # chosen sustained useful rates and group prices
2 {"name": "two-GPU group", "tps": 620, "hourly_cost": 6.0},
3 {"name": "four-GPU group", "tps": 900, "hourly_cost": 12.0},
4]
5for candidate in eligible:
6 candidate["cost_per_million"] = candidate["hourly_cost"] / (candidate["tps"] * 3600) * 1_000_000
7 print(f"{candidate['name']}: ${candidate['cost_per_million']:.2f}/million useful tokens")
8fastest = max(eligible, key=lambda candidate: candidate["tps"])
9cheapest = min(eligible, key=lambda candidate: candidate["cost_per_million"])
10print(f"highest rate: {fastest['name']}; lowest GPU-only unit cost: {cheapest['name']}")
11assert fastest is not cheapest1two-GPU group: $2.69/million useful tokens
2four-GPU group: $3.70/million useful tokens
3highest rate: four-GPU group; lowest GPU-only unit cost: two-GPU groupThe prices and rates are arithmetic assumptions, not provider quotes or a benchmark. Whole-service cost also includes idle intervals, other workers, retries, and CPU/network/storage charges.
How do you choose between model parallelism and replicas?
Answer
Use model parallelism when a single model copy exceeds single-GPU VRAM or a single request requires multi-GPU compute. Use replicas when the model fits and request volume is the bottleneck. Use both when serving large models under heavy traffic: shard the weights across minimal GPUs, then replicate that sharded group to scale throughput.
When multi-GPU serving breaks down
These symptoms guide investigation without identifying a cause by themselves:
- OOM despite on-paper fit: Inspect each rank's checkpoint/replication, KV layout, workspace, graph capture, other allocations, and reported free bytes. A group-total sum can hide a local deficit.
- Latency worsens with more ranks: Profile collectives, local kernel efficiency, scheduling, and topology. Test whether exposed communication exceeds saved work; a PCIe or cross-node label alone doesn't prove dominance.
- KV history stops shrinking with TP: Check head replication and hybrid-layer state. Weight savings can still enlarge the pool. Test supported DCP if history duplication is the measured problem.
- Concurrent traffic creates token gaps: Separate queued admission, mixed-step work, transfers, and delivery buffering in a trace. Chunking or phase separation may help; measure their added work and queues too.
Evaluation rubric
- Reconstruct the dense MLP result from column-parallel and row-parallel shards, adding output bias once.
- Count collectives under the classic dense layout and measure which communication is exposed.
- Derive the pipeline bubble fraction from first principles.
- Distinguish Ring Attention from DeepSpeed Ulysses sequence sharding for long-context workloads.
- Filter by per-rank fit and the actual latency contract, then compare sustained useful rates and billed costs.
Follow-up questions
What changes if all twelve expert assignments land on device 0?
Answer
The peak-to-average assignment ratio is 12/6 = 2.0 across two devices. The fixture's weighted output remains correct if every assignment is computed. Counts expose skew, but expert/kernel costs, local work, and transfers are needed to predict the straggler or elapsed latency.
Why doesn't 58 GB of free group memory prove that a 50 GB request fits?
Answer
Placement is local. A request needing 35 GB on rank 0 and 15 GB on rank 1 exceeds rank 0's 29 GB remainder even when their sum is 58 GB. That layout can't be admitted as-is; waiting, rejection, or preemption depends on policy. Every rank must satisfy its allocation.