Model Parallelism for LLM Inference

HBM, KV cache, and scheduler policy limit single-node serving. The next question is what changes when a single large language model (LLM) copy no longer fits comfortably on one accelerator.

Serving Qwen3.6-35B-A3B for a codebase assistant that reads long diffs, build logs, and architecture notes creates two memory questions. Model-weight memory decides whether one replica fits at all, while per-request KV state decides how many long-context sessions can run at once.

Qwen3.6-35B-A3B is a sparse MoE checkpoint with about 35B total parameters and about 3B activated per token.^[1] At BF16, the full checkpoint is roughly 70 GB before KV cache, runtime buffers, and allocator headroom. The A3B suffix helps reason about active compute, but it doesn't mean the serving system only needs 3B parameters worth of memory.

Model parallelism is the set of techniques that split one model across multiple accelerators. The split can make a large model fit, but every added device also creates communication or scheduling work.

Why inference sharding differs from training

Distributed training cares about gradients, optimizer states, activation checkpointing, and throughput over many examples. Distributed inference cares about time to first token (TTFT), tokens per second (TPS), KV-cache memory, and request scheduling.

The same names appear in both worlds, but the trade-offs shift:

Serving runtimes expose topology controls such as tensor- and pipeline-parallel sizes. Treat those knobs as a deployment mechanism, not a performance guarantee: the selected topology still needs memory accounting and latency benchmarks. vLLM's official scaling guide recommends one GPU when the model fits, single-node tensor parallelism when it needs several GPUs in one node, and tensor plus pipeline parallelism when it exceeds one node.^[2]

Before picking a sharding strategy, do the simplest memory math:

For BF16 or FP16 weights, a parameter takes 2 bytes. That makes Qwen3.6-35B-A3B about 70 GB by total parameters before a single prompt arrives. KV cache grows with active requests, context length, layers, KV heads, head dimension, and bytes per value. Runtime buffers and memory fragmentation add more headroom. "Enough VRAM" means all of those buckets fit at the traffic level you plan to serve, not the checkpoint file or active-parameter count alone.

This calculator is deliberately conservative but incomplete. It only sizes weights plus a reserve. Production sizing still needs KV cache, activation buffers, tensor-parallel divisibility, interconnect measurements, and latency SLOs.

Tensor parallelism

Tensor parallelism splits large matrix operations across GPUs. In a transformer layer, the model contains big linear projections for attention and feed-forward blocks. Tensor parallelism shards those weights so each GPU owns part of the matrix.

For a simplified linear layer:

you might split W across four GPUs. Each GPU computes part of the output, then the system communicates to combine results. Interconnect decides whether that split pays off. Tensor parallelism is strongest on GPUs connected by fast NVLink (NVIDIA high-bandwidth GPU-to-GPU interconnect) or similar high-bandwidth links.

The Megatron pattern: column then row

Naively, every sharded matmul would need a sync to glue the pieces back together. Megatron-LM avoids that by choosing the split directions so each transformer sub-block needs only one synchronization.^[3] It chains a column-parallel layer into a row-parallel layer.

In the MLP block Z = (GeLU(xA))B:

• Split the first weight A by columns. Each GPU computes GeLU(x A_i) independently. GeLU is element-wise, so no sync is needed before the nonlinearity. (Splitting A by rows would force a sync first, because GeLU(x1 A1 + x2 A2) isn't the sum of the per-shard GeLUs.)
• Split the second weight B by rows. The column-sharded output of the first layer is exactly the input layout the row-parallel second layer expects, so the partial results flow through with no intermediate communication. One all-reduce after B sums the partial outputs.

The attention block follows the same shape. The query, key, and value projections are split column-wise, which maps cleanly onto independent attention heads, and the output projection is split row-wise. That gives one all-reduce after attention.

For the dense Megatron layout described here, a tensor-parallel group executes two all-reduces per transformer layer in the forward pass: one for attention and one for the MLP.^[3] During decode, that means every generated token pays two collectives per layer. An 80-layer model therefore executes 160 all-reduces per decode step under this layout, unless a runtime changes or fuses the communication scheme.

For inference, tensor parallelism often reduces memory pressure and can improve latency for large models, but it adds communication inside layers. If communication is slow, adding GPUs can make serving worse.

The interconnect hierarchy is why tensor parallelism is usually easiest to justify within a fast-link domain. NVIDIA's H100 SXM specification advertises up to 900 GB/s NVLink aggregate bandwidth per GPU.^[4] PCIe and cross-node network paths have different bandwidth and latency characteristics, so the same tensor-parallel layout can perform very differently across machines. Because every decode token triggers collectives, benchmark the target topology instead of assuming more GPUs lower latency.

Use tensor parallelism when:

1. The model doesn't fit on one GPU.
2. The GPUs have fast interconnect.
3. You need one request to use multiple GPUs at once.
4. Batch sizes aren't large enough to rely only on model replicas.

Megatron-LM popularized practical tensor model parallelism for large transformers.^[3] In serving stacks, the same idea appears as tensor_parallel_size.

Pipeline parallelism

Pipeline parallelism splits layers into stages. GPU 0 owns early layers, GPU 1 owns middle layers, and GPU 2 owns later layers. A token's hidden state moves through the stages.

This reduces memory per GPU because each stage stores only part of the model. It can also reduce communication compared with tensor parallelism because tensors move between stages rather than across every large matmul.

Pipeline parallelism creates bubbles. If only one request is active, stage 2 waits for stage 1, stage 3 waits for stage 2, and so on. Bigger batches or many concurrent requests can fill the pipeline better.

Use pipeline parallelism when:

1. Tensor parallelism alone doesn't fit the model.
2. The model must cross node boundaries.
3. You can batch enough work to keep stages busy.
4. You can tolerate slightly more scheduling complexity.

Tensor and pipeline parallelism can be combined. For example, eight GPUs can run tensor_parallel_size=4 and pipeline_parallel_size=2, giving two layer stages where each stage is a four-GPU tensor-parallel group. For a multi-node vLLM deployment, the common first layout is tensor parallelism inside each node and pipeline parallelism across nodes. vLLM also recommends considering pipeline parallelism inside one node when GPU count doesn't evenly divide the model or the node lacks NVLink.^[2] This may make a larger replica fit, but queue depth and interconnect measurements still decide TTFT and throughput.

Sequence parallelism and context parallelism aren't interchangeable

Both names mention the token dimension, but they solve different problems. In Megatron Core, sequence parallelism works alongside tensor parallelism: it shards sequence-dimension work in components such as LayerNorm and Dropout to reduce activation memory. Context parallelism partitions the sequence across devices through the transformer layers and is the long-sequence strategy in Megatron's current guide.^[5]

For inference, support depends on the runtime and model architecture. Context sharding means the system distributes long-prompt attention work or KV state across devices instead of only splitting weights. Ring Attention is one family of techniques for this problem. Devices hold local query blocks while Key and Value blocks circulate through a ring for blockwise attention. The paper applies this idea to training and inference.^[6]

For a codebase assistant reading a long repository map plus build logs, sequence length can dominate prefill cost. Tensor parallelism helps with model weights. Prefix caching helps with repeated prefixes. Context-aware serving helps when the prompt itself is large and attention work or KV state needs to be spread out.

Context parallelism isn't the first knob most teams touch. Start with model size, quantization, tensor parallelism, and batching. Reach for context-level techniques when long prompts are the bottleneck and your runtime supports the required communication pattern.

Expert parallelism for MoE models

Mixture-of-Experts models add another sharding axis: experts. Instead of every token using every feed-forward block, the router sends each token to a small subset of experts.^[7] Expert parallelism places different experts on different GPUs, so the serving system can scale total expert capacity without copying every expert to every device.

Expert parallelism pays in routing communication and load balance. Expert-parallel implementations commonly dispatch tokens to devices that own selected experts and combine results afterward, often using all-to-all-style communication. If many tokens choose the same expert, that expert's device becomes the bottleneck while other devices wait. For dense models, start with tensor/pipeline/context choices. For MoE serving, add expert placement and router-load metrics to the plan.

DeepSeek-V3, for example, reports 671B total parameters with 37B activated per token, illustrating why total expert storage and active-token compute are different capacity questions.^[8] Serving such a model still needs expert placement, routing balance, and communication measurements; expert parallelism can combine with tensor and data parallelism rather than replace them.

Sizing example

Suppose you need to serve Qwen3.6-35B-A3B for a codebase-reasoning assistant:

A reasonable benchmark candidate is one fast-linked node with enough high-memory GPUs for one replica, tensor parallelism within that node, continuous batching, and prefix caching for stable policy text. If its cost or latency is unacceptable, measure quantization before adding cross-node parallelism.

Some systems add one more axis: disaggregated serving runs prefill and decode on separate worker pools, each with its own parallelism, and transfers KV state between them. Systems such as DistServe and Splitwise study when this can raise goodput by reducing phase interference, subject to KV-transfer overhead.^[9][10] The parallelism choices below still apply, but they can be evaluated per phase.

What to measure

Multi-GPU inference should be measured with serving metrics, not offline tokens per second alone.

Track:

1. Time to first token.
2. Decode tokens per second.
3. Aggregate throughput.
4. GPU memory used by weights.
5. GPU memory used by KV cache.
6. Interconnect utilization.
7. Queue time under burst traffic.
8. Error rate when one GPU or node fails.

The worst mistake is counting total VRAM and declaring victory. A four-GPU box with enough raw memory can still miss latency targets if the interconnect is saturated or the scheduler can't fill the pipeline.

Model parallelism is a capacity tool for models that need multiple GPUs. Replicas fit many independent requests when each model copy fits on one GPU. Combine both when the product needs a large model and real throughput.

Practice: choose sharding or replicas

Consider two serving plans:

This is the decision habit to build: ask whether the bottleneck is model weight memory, request volume, context length, or communication. Different bottlenecks need different tools.

Mastery check

• Why a model may need sharding for inference even after quantization.
• How tensor parallelism, pipeline parallelism, sequence parallelism, context parallelism, expert parallelism, and replicas solve different bottlenecks.
• Why sequence parallelism and context parallelism aren't interchangeable names for long-context serving.
• Why the Megatron column-then-row pattern needs only two all-reduces per layer, and what that costs on every decode token.
• Why tensor parallelism depends on fast interconnect (NVLink vs PCIe vs network) and can become communication-bound.
• Why pipeline parallelism can save memory but hurt small-batch latency through bubbles.
• How disaggregated prefill and decode let each phase pick its own parallelism layout.
• How to design a first serving plan from memory buckets, traffic shape, context length, and latency targets.

Evaluation rubric

Strong answers should:

• identify the real bottleneck before recommending replicas or sharding
• separate memory-fit math from latency and communication measurements
• explain why tensor, pipeline, context, and expert parallelism solve different limits
• connect interconnect speed directly to TTFT and decode behavior
• name the smallest production-worthy first plan instead of the fanciest one

Follow-up questions

Common pitfalls

• Symptom: The model fits on paper, but the runtime still hits OOM. Cause: You counted weight memory and forgot KV cache, runtime buffers, allocator slack, and bursty long-context headroom. Fix: Size memory by bucket, not checkpoint size alone.

• Symptom: Latency gets worse after you spread the model across more GPUs. Cause: Cross-node tensor-parallel collectives now dominate decode. Fix: Keep the tensor-parallel group inside one fast node first. Only cross weaker links when pipeline behavior and queue depth justify it.

• Symptom: TTFT rises even though the bigger shard plan finally fits. Cause: Extra communication and startup coordination removed less pressure than they added. Fix: Measure TTFT and decode TPS directly. More GPUs aren't automatically a serving win.

• Symptom: You shard a model that already fits, but throughput barely improves. Cause: Traffic volume was the real bottleneck, so communication replaced a simpler replica plan. Fix: Use replicas first when one full model copy fits and requests are independent.