@_avichawla: https://x.com/_avichawla/status/2107427400779026843
Summary
An in-depth explainer covering parallelism strategies for LLM inference—data, tensor, pipeline, context, and expert parallelism—explaining how to split inference across GPUs without losing performance to communication overhead. It walks through prefill vs. decode phases, memory and communication bottlenecks, and how to choose a defensible layout given model, hardware, and workload constraints.
View Cached Full Text
Cached at: 10/06/26, 05:54 PM
Parallelism strategies for LLM inference, clearly explained
Everything you need to understand how to split LLM inference across GPUs without losing performance to communication. It covers data, tensor, pipeline, context, and expert parallelism, how to combine them, and the deployment checks that determine which layout fits your model, hardware, and workload.
Four GPUs provide four separate memory spaces and four processors. They do not behave like one GPU with four times the capacity.
A kernel running on one GPU cannot use weights or KV-cache blocks stored on another GPU without transferring that data.
Distributed inference therefore starts with data placement. The serving engine assigns weights, activations, requests, and cached keys and values to specific devices.
Any dependency that crosses devices adds a transfer or collective operation to the request path.
The engineering problem is choosing a placement that removes the current limit without adding more communication than the hardware can absorb. That is what we will solve here.
By the end, you should be able to inspect a model, workload, GPU topology, and latency target, then select a defensible first layout.
We will compare each strategy with five questions:
-
What state is partitioned?
-
What remains replicated?
-
What crosses the interconnect?
-
How often does that transfer occur?
-
Does the layout improve memory capacity, request latency, or throughput?
Throughout the article, we will assume the same four-GPU server. The model needs at least two of those GPUs, and the server receives both short chats and long-context requests. We will keep that model, workload, and hardware fixed while changing only how inference work is divided across the GPUs.
We will use the word “collective” throughout this article. A collective is an operation in which several GPUs exchange or combine data as a group. For example, two GPUs may calculate separate parts of a layer output. An all-reduce adds those parts together and gives the complete result to both GPUs. Other collectives move different pieces of data between the participating GPUs. Every participating GPU must complete its part before the collective finishes. One slow GPU or connection can therefore delay the entire group. This differs from a direct transfer, where one GPU sends data to another GPU.
How LLM inference works
Before we can understand how LLM inference can be optimized, we need to see how LLM inference works in the first place.
LLM inference has two phases.
-
Prefill processes the prompt and builds its key-value cache. It handles many prompt tokens together and often uses large matrix multiplications. Prefill strongly affects time to first token, usually called TTFT.
-
Decode generates later tokens one step at a time. Each step reads model weights and the growing cache. Decode strongly affects time per output token, usually called TPOT. Some systems report inter-token latency instead.
These phases use hardware differently.
-
Prefill usually executes large matrix multiplications across many prompt tokens.
-
Decode executes smaller operations and repeatedly reads weights and KV-cache data. A layout that shortens prefill can still slow decode by adding synchronization to every generated token.
I have covered this in detail here if you want to dive deeper:
Avi Chawla@_avichawla·Jun 28 ArticleHow LLM Inference Works, Clearly Explained.Every generate() call to an LLM runs two distinct computational phases on the same GPU: prefill (processing the prompt) is compute-bound while decode (generating tokens one at a time) is memory-bound….10105764308K
Four resources usually limit the deployment.
-
Parameter memory determines whether the model fits.
-
KV-cache memory limits resident tokens and concurrent requests.
-
GPU compute time covers local matrix operations.
-
Communication time covers transfers and collectives between devices.
The service metrics must match the underlying product target.
-
TTFT measures initial responsiveness.
-
TPOT measures streaming pace.
-
Request throughput counts completed requests.
-
Token throughput separates input and output work.
-
Goodput counts only requests meeting a service-level objective.
Considering averages across any metric will inherently hide tail latency. That is why one must report the median and upper percentiles.
In addition to this, the p95 TTFT is the threshold that 95 percent of requests meet. Compare it with a service objective under the same request distribution.
Note. Tokens per second without a latency limit can be misleading. A saturated server may report high throughput while violating its tail-latency objective.
These metrics tell us which service objective is failing. And the resource measurements tell us whether memory, compute, or communication caused it.
The next decision is how to divide inference work across the GPUs to address that specific limit.
That is the role of parallelism. Each method partitions a different part of the workload and introduces a different communication workload to be handled. The comparison below provides a map of the techniques we will examine next.
The communication column above is very important. While the transfer size itself is important, the frequency of transfer does too. For instance, a small collective repeated in every decoder layer can cost more than one larger activation transfer between pipeline stages.
We will begin with data parallelism because it keeps each forward pass inside one replica. The other methods split a single forward pass across GPUs and therefore place communication directly on its execution path.
1) Data parallelism
Data parallel inference partitions requests.
Each replica stores the complete model. It also owns an independent scheduler and KV cache. A router sends each request to one replica.
This differs from the training-time data parallelism.
Since the inference phase has no requirement for gradient synchronization after each forward pass, one request remains local to its chosen replica.
Suppose one two-GPU replica meets the latency target but cannot handle the incoming request rate.
On our four-GPU server, data parallelism lets us run a second two-GPU replica. The replicas process different requests concurrently, so this configuration targets throughput rather than single-request latency.
- The first deployment check is memory.
Each replica must contain the complete two-GPU model layout, which means the four-GPU server stores the model twice.
Data parallelism is therefore available only after one complete replica already fits. It cannot solve a model-placement problem by itself.
- The second check is request routing.
Once two replicas are active, the router decides which queue and KV-cache pool receives each request.
Round-robin routing treats both replicas as equally available, even when one has a longer queue, less free KV capacity, or a reusable prefix already cached.
- The third check is admission capacity.
A request consumes little KV memory when it arrives, then consumes more as decode adds tokens.
The router must reserve capacity for that growth. Otherwise, a replica can accept a request that fits during prefill but runs out of KV space during generation.
Here’s an engineering rule for this. After adding replicas, route with queue delay, available KV capacity, and prefix-cache hits. Reserve KV space for the expected decode length before admitting the request.
This also explains when data parallelism does not help.
At low traffic, the second replica remains idle and cannot reduce the latency of one request. Use it when a complete replica fits, and request concurrency is high enough to keep multiple replicas busy.
2) Tensor parallelism
Suppose the model does not fit on one GPU, but does fit when its largest weight matrices are divided across two.
Data parallelism cannot solve this problem because every replica still needs a complete model layout. Tensor parallelism can because it makes several GPUs cooperate on the same transformer layer and the same request.
This makes tensor parallelism a model-placement tool first. It can also reduce the compute time of a large operation, but only when the saved compute exceeds the cost of communicating between GPUs.
The original Megatron-LM paper explains the common two-projection construction. We can see the dependency through one feed-forward block.
The first projection creates an intermediate activation. The second projection returns it to the hidden width.
The following diagram labels the input, both weight matrices, the activation function, the intermediate tensor, and the output.
With two ranks, we split the first weight matrix by columns. Each rank computes different intermediate features.
We split the second matrix by rows, so each rank consumes its matching feature slice. Both ranks then produce partial output tensors.
Each GPU computes one matrix product with local weight shards. The all-reduce sums those partial tensors. Neither GPU has the correct layer output before that collective completes.
pythonimport numpy as np
def tensor_parallel_output(x, a_shards, b_shards, activation=np.tanh):
partials = []
for a_shard, b_shard in zip(a_shards, b_shards, strict=True):
hidden_shard = activation(x @ a_shard)
partials.append(hidden_shard @ b_shard)
return np.sum(partials, axis=0)
The function follows the two-rank data dependency. Each rank uses matching matrix shards. The reduction is required because every local result contains only part of the output.
Running the tested example produces this output:
plaintexttensor_parallel_output = [[-3.294438, 1.406896]]
The unsharded MLP produces the same two values. The test checks every output element within NumPy’s floating-point tolerance.
Attention uses the same broad idea.
-
Query, key, and value projections can split across heads.
-
The output projection then combines partial results.
-
Exact layouts depend on multi-head attention, grouped-query attention, and model-specific kernels.
Before choosing a tensor-parallel size, two terms need to be clear.
-
A tensor-parallel group is the set of GPUs serving one request together.
-
A rank is one GPU process inside that group. TP2 has two ranks, while TP4 has four.
Increasing the TP degree divides each sharded matrix into more pieces. Consider a matrix with 8,192 output columns.
-
TP2 gives 4,096 columns to each rank.
-
TP4 assigns 2,048 columns to each rank.
So each GPU stores fewer weights and performs a smaller matrix multiplication.
That sounds like a direct improvement, but the layer cannot continue with separate partial results.
The ranks must exchange or combine them.
TP4 therefore performs less local work per rank while involving twice as many ranks as TP2. Whether TP4 is faster depends on the balance between saved computation and added communication.
This gives us three deployment checks.
- The first check is whether the model actually fits.
Sharding reduces memory for the divided weight matrices. Other allocations may remain replicated on every rank.
These include some parameters, temporary buffers, CUDA graphs, and KV-cache allocations. Measure peak memory on each GPU. Do not assume TP2 halves every allocation.
- The second check is how long ranks spend communicating.
Tensor-parallel collectives combine partial results inside many transformer layers. Slow links can leave GPUs waiting after each local calculation. This wait occurs during prefill and every decode step.
- The third check is whether each rank still receives enough work.
GPUs run large matrix multiplications efficiently because they inherently expose substantial parallel work.
Raising the TP degree makes each local matrix smaller. At some point, the smaller calculation uses the GPU less efficiently, while the collective still occurs.
Attention heads and quantization layouts impose additional constraints. For example, a head or quantization group may need to remain intact on one rank. The model shape and serving engine therefore determine which TP degrees are valid.
TP2 can outperform one GPU on a large prefill because both ranks receive substantial work.
TP4 may help further when compute still dominates. It may also be slower when smaller local operations no longer cover the added communication.
Decode reaches this point sooner because one token provides less matrix work than a long prompt.
One common engineering rule is to use the smallest tensor-parallel degree that meets memory and latency requirements. Measure again after every increase.
The links between ranks determine the communication cost. NVLink and NVSwitch differ from an inter-node network in latency and bandwidth.
For our four-GPU server, start with TP2 because the model needs two GPUs.
If that layout meets the latency target, the remaining two GPUs can form a second TP2 replica for throughput.
Move to TP4 only if TP2 still fails a memory or single-request latency requirement. TP4 spends all four GPUs on each request and removes the option of running a second replica.
Use tensor parallelism to solve a per-replica fit problem or a measured single-request compute limit. Do not increase the degree merely because more GPUs are available.
3) Pipeline parallelism
Tensor parallelism is expensive when its per-layer collectives cross a slow link.
If the model must span GPUs with weaker connectivity, pipeline parallelism offers another placement, i.e., divide the model by depth, keep a contiguous group of layers on each stage, and transfer activations only when execution moves between stages.
So stage zero runs the first layer group and sends its hidden states to stage one. The final stage runs the last group and produces logits.
The request still visits every stage in order, but the weights for any one stage remain local to its GPU or tensor-parallel group.
This differs from the traditional training pipeline.
Inference has no backward pass. Continuous batching can keep several request groups moving. One autoregressive sequence still depends on its preceding token.
Pipeline parallelism therefore solves a model-placement problem with less frequent communication than tensor parallelism.
That makes it a useful candidate across nodes or GPUs without fast peer-to-peer links. It does not automatically reduce the latency of one request, because that request must still pass through all stages.
Microbatching supplies independent work to the stages.
During online serving, those microbatches may contain different requests. Continuous batching changes membership as requests finish. The scheduler must preserve pipeline order while filling newly available slots.
The first deployment check is whether the workload supplies enough independent requests to keep the stages busy.
Stages are idle during warm-up and drain. Too few microbatches increase this idle fraction, so a lightly loaded online service may use only one stage at a time.
The second check is stage balance.
An 80-layer model suggests 20 layers per stage in PP4, but equal layer counts do not guarantee equal execution time or memory use. Embeddings, output layers, and heterogeneous blocks distort the split. The slowest stage sets throughput while faster stages wait.
One engineering rule is to balance stages with measured execution time and memory use. Equal layer counts are only an initial placement.
The third check is activation buffering and operational cost. Activations remain live while requests move through the pipeline. The scheduler must maintain order, and a failure on one stage affects the whole model instance.
For our four-GPU server, PP4 is reasonable only if the model needs all four GPUs and enough requests are in flight to fill four stages.
If the hardware has two fast GPU pairs connected by a slower path, PP2 with TP2 inside each stage can keep frequent tensor collectives within each pair and send activations across the slower boundary.
Use pipeline parallelism when model depth must span devices and its stage-boundary transfers fit the topology better than per-layer tensor collectives. Avoid it for low-concurrency traffic unless model fit leaves no alternative.
4) Context parallelism
Assume the weights now fit, but long-context requests still fail one of two tests.
A large prompt may miss the TTFT target during prefill, or its KV cache may consume so much memory that too few requests can remain active during decode.
Context parallelism addresses these problems by partitioning sequence positions rather than model weights.
The two failures require different mechanisms.
-
Prefill context parallelism, or PCP, divides the work for an incoming prompt.
-
Decode context parallelism, or DCP, divides the stored KV history used by later tokens.
Prefill context parallelism
Prefill context parallelism divides prompt tokens across ranks. Each rank computes query, key, and value tensors for its local token chunk. The attention algorithm must still expose every query to the required keys and values.
One implementation gathers keys and values on every rank. Each rank evaluates its local queries against the complete history.
This works only while the gathered tensors fit in memory.
Another family circulates partial key-value chunks. Ring-attention-style algorithms merge partial attention results incrementally.
They avoid materializing the full sequence everywhere, but require several communication rounds.
The practical question is where PCP crosses over from overhead to benefit. Measure TTFT across prompt-length buckets, not only at the maximum context length.
Long prompts may contain enough attention work to offset the exchanges. Short prompts usually do not, so a mixed service may route only long-context requests to PCP-enabled workers.
Decode context parallelism
Decode context parallelism partitions the stored KV cache by sequence position.
Each rank attends to its local history. The ranks then merge partial attention outputs and softmax statistics.
pythonimport numpy as np
def merge_attention_partials(partials): global_max = max(part.maximum for part in partials) scales = [np.exp(part.maximum - global_max) for part in partials] denominator = sum( scale * part.exponential_sum for scale, part in zip(scales, partials, strict=True) ) numerator = sum( ( scale * part.weighted_value for scale, part in zip(scales, partials, strict=True) ), start=np.zeros_like(partials[0].weighted_value), ) return numerator / denominator
Each rank reports its local maximum, exponential sum, and weighted-value sum.
The global maximum rescales every contribution before the final reductions. This preserves stable softmax across KV shards. The project file tests the result against unsharded attention.
The local run produces this merged attention output:
plaintextmerged_attention_output = [0.769627, 0.779901]
The full attention calculation produces the same result. This verifies that splitting the KV history does not change the attention output for the test input.
The maximum is needed because exponentials can overflow. Subtracting the global maximum keeps every exponent at most one.
The scaled denominators and weighted-value sums then reconstruct the unsharded attention result within normal floating-point tolerance.
- The first DCP check is memory.
Record resident KV tokens and maximum concurrency before and after sharding.
- The second is decode cost.
Every decoder layer must merge partial attention results, which can increase TPOT. DCP is useful only when the extra KV capacity or supported context length is worth that repeated merge.
Sequence parallelism is related, but not identical.
Megatron-style sequence parallelism commonly shards activations around normalization and similar operations. Context parallelism partitions attention across context positions. Frameworks sometimes use these labels differently.
For our four-GPU server, keep ordinary chats on the baseline TP2 replicas. Route the occasional long prompt to a PCP-enabled group if prompt-length tests show a TTFT improvement.
Add DCP only if KV memory, rather than weight memory or compute, limits the number or length of active requests.
PCP solves long-prompt compute. DCP solves long-history storage. Treat them as separate decisions and enable neither without a prompt-length or KV-capacity measurement that justifies it.
5) Expert parallelism
Expert parallelism applies only when the model contains mixture-of-experts layers.
In that case, storing every expert on every GPU can dominate parameter memory even though each token uses only a few experts.
Expert parallelism places different experts on different ranks so one model instance can distribute that expert-weight memory.
At runtime, the router chooses experts for each token. The system sends each token representation to the ranks that own those experts, runs the local expert networks, and returns their outputs.
The placement saves memory, but it puts data-dependent all-to-all traffic on the request path.
Avi Chawla@_avichawla·Sep 18 ArticleMoE inference engineering, clearly explainedEverything you need to understand how MoE serving engines route tokens, batch expert computation, distribute weights across GPUs, and manage communication and load imbalance. It covers expert…849316147K
The first deployment check is placement.
Count expert-weight memory per rank and verify that shared attention, embeddings, buffers, and KV cache still fit. Expert parallelism does not determine how the dense parts of the model are placed.
The second check is routing traffic. Unlike dense tensor parallelism, the transferred bytes depend on which experts the router selects. Measure bytes sent per rank and all-to-all time using the real prompt distribution.
The third check is imbalance. Token counts can skew across experts, and the collective finishes when the busiest rank finishes.
Average load hides this straggler. Measure the maximum tokens per expert, local expert time, and post-compute wait time. Expert capacity, padding, and load-balancing policy can change both runtime and, if tokens are dropped, model behavior.
On four GPUs, an MoE model might distribute one quarter of its experts to each GPU while using TP or DP for attention and other dense operations.
This layout is attractive when expert weights are the fit problem and routing remains balanced. It is unattractive when the expert set already fits comfortably or all-to-all communication dominates layer time.
Attention and expert layers need not use the same partition. Dense attention may use DP or TP while MoE layers use EP.
Use expert parallelism to solve expert-weight placement. Validate it with per-rank routing and imbalance measurements, not with average token throughput alone.
6) Hybrid parallelism
One parallel method often solves only part of the deployment.
Tensor parallelism may make a model fit within a node, pipeline parallelism may extend it across nodes, and data parallelism may then add enough replicas for the request rate.
A hybrid layout composes those choices, but each added axis also adds its own scheduler or communication path.
For independent dense-model groups, the worker count is the product of the data, pipeline, and tensor-parallel degrees.
Some context-parallel methods reuse existing ranks. Expert groups can also interact with data groups differently.
Consider eight GPUs across two four-GPU nodes. The intra-node links have higher bandwidth and lower latency. One candidate layout uses TP4 inside each node and PP2 across nodes.
This keeps frequent layer collectives local. Only stage-boundary activations cross the slower network. All eight GPUs extend one model instance.
If the model fits on four GPUs, another layout becomes attractive. TP4 within each node and DP2 across nodes creates two replicas. Requests no longer depend on the cross-node path.
Build a hybrid layout in that order. First choose the smallest TP group that makes each stage fit, and keep it on the fastest links.
Add PP only when the model must cross a slower boundary. Add DP after one complete model instance fits. Add CP for a measured long-context problem and EP only for MoE expert placement.
This order prevents a common mistake: multiplying parallel degrees before deciding which resource each one fixes. Every new axis should have one stated job and one measurement that verifies it.
Prefill and decode disaggregation is useful when the two phases need different hardware or scaling ratios.
Separate worker pools handle each phase, so the prefill pool can favor large prompt operations while the decode pool favors KV capacity and steady token generation. This introduces a KV-state transfer and coordination boundary.
Disaggregation does not replace TP, PP, or CP. Each worker pool still needs a model-placement plan.
A hybrid layout is successful only when each communication pattern matches the physical network. Place the most frequent collectives on the fastest links, then verify the complete layout under the production request mix.
A practical decision framework
Use this framework from top to bottom.
Stop after selecting the smallest change that addresses the measured limit, benchmark it, and return to the framework only if another limit remains.
Stacking every available parallel mode makes attribution harder and can add communication without solving a problem.
Start by determining whether one complete replica fits on one GPU while preserving KV-cache capacity for the target concurrency.
-
If it fits, start with one replica. Add data parallelism when traffic can use another scheduler. Preserve prefix locality in the router.
-
If it does not fit, find the smallest tensor-parallel degree that works. Benchmark decode latency at that degree. Consider pipeline parallelism when layer groups have similar execution time or peer links are slow.
If the model must span nodes, keep TP within the faster intra-node links. Test PP across nodes. Measure stage balance before increasing load.
Long prompts create separate questions. Evaluate PCP when TTFT fails at large input lengths. Evaluate DCP when KV capacity limits concurrency or long-context decode.
MoE models need one more branch. If expert weights cause the fit problem, test EP while tracking token imbalance and all-to-all time separately. Do not add EP to a dense model or infer its benefit from total GPU utilization.
Enable a parallel mode only when measurements identify the resource limit it addresses.
Identify the constrained resource first. Then choose a partition whose communication fits the available interconnect.
Parallelism is typically a placement decision, not a GPU-count setting. Start with the resource that prevents the service from meeting its target. Then choose the smallest layout change that addresses that constraint.
Every split introduces communication, synchronization, or scheduling overhead. The saved memory or compute must outweigh that cost on the actual hardware and request distribution.
The best layout is the simplest one that meets the latency and throughput objectives. Everything else still needs measurement.
That’s a wrap!
If you enjoyed this:
Find me → @_avichawla
Every day, I share tutorials and insights on DS, ML, LLMs, and RAGs.
Similar Articles
@_avichawla: Prefill & decode in LLM inference. Have you ever noticed that the first token from an LLM always takes a moment to appe…
Explains the two phases of LLM inference - prefill and decode - detailing how GPU bottlenecks shift from compute-bound during prefill to memory-bound during decode, and the importance of KV caching.
@_avichawla: The LLM scheduling loop, clearly explained: (bookmark this) Whenever a user sends a prompt to an LLM, the resulting inf…
An educational walkthrough of the LLM serving scheduling loop, explaining how inference requests move through waiting queues, batch construction, prefill, decode, and continuous batching to keep GPUs saturated as requests finish at different times.
@akshay_pachaar: https://x.com/akshay_pachaar/status/2087928032904523980
An educational thread explaining how GPUs work, focusing on the memory-compute asymmetry that dominates LLM serving performance, and demonstrating how techniques like quantization, speculative decoding, and continuous batching follow from that fundamental constraint.
An open handbook on LLM inference at scale (GPU internals, KV cache, batching, vLLM/SGLang/TensorRT-LLM) [P]
An open, in-progress handbook explaining LLM inference internals including GPU memory hierarchy, KV cache, batching, and popular inference engines like vLLM and TensorRT-LLM.
@_avichawla: https://x.com/_avichawla/status/2077653695123378321
This article argues that vLLM and similar serving frameworks are inefficient for running multiple small AI models on a single GPU due to design limitations. It introduces the SIE open-source inference engine as a solution for serving many models together to reduce costs.