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]Reference 1Qwen3.6-35B-A3Bhttps://huggingface.co/Qwen/Qwen3.6-35B-A3B} 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.

Current open-weight models span very different sharding problems:

GLM-5.2 and Flash 0731 are not bigger versions of the Qwen example. Both combine sparse expert routing with specialized long-context attention, so a serving plan must check runtime support for the exact architecture and checkpoint. Active parameters predict neither full weight residency nor collective cost.^{[2]Reference 2DeepSeek-V4-Flash-0731 Model Cardhttps://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731[3]Reference 3DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligencehttps://arxiv.org/abs/2606.19348[5]Reference 5GLM-5.2 Model Cardhttps://huggingface.co/zai-org/GLM-5.2[6]Reference 6GLM-5 Repositoryhttps://github.com/zai-org/GLM-5[7]Reference 7GLM-5.2 Configurationhttps://huggingface.co/zai-org/GLM-5.2/blob/main/config.json}

Why can Qwen3.6-35B-A3B need more than one GPU before serving any real traffic?

Answer

The BF16 checkpoint is roughly 70 GB, while an 80 GB GPU with a conservative 20% reserve has only about 64 GB usable. Serving also needs KV cache, runtime buffers, fragmentation, batching, and safety headroom, so active-parameter count is not a memory budget.

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.^{[8]Reference 8Distributed Inference and Serving.https://docs.vllm.ai/en/latest/serving/parallelism_scaling.html}

Why can't you reuse the same mental model for distributed training and distributed inference?

Answer

Training optimizes gradient throughput and optimizer-state memory. Inference optimizes TTFT, decode TPS, KV-cache capacity, scheduler behavior, and communication cost on every generated token.

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.^{[9]Reference 9Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.https://arxiv.org/abs/1909.08053} 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.

Head-count divisibility. That head-parallel story assumes the tensor-parallel size divides the head counts. In practice TP must divide num_attention_heads, and under GQA the tighter constraint is often num_kv_heads (TP ≤ kv heads, or the runtime pads / re-layouts). Many open MoE and hybrid models break naive head-parallel TP at large TP sizes; load errors and uneven shards show up at deploy time, not in the Megatron textbook diagram.

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.^{[9]Reference 9Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.https://arxiv.org/abs/1909.08053} 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.

In a tensor-parallel linear layer, what does each GPU compute and why is communication required afterward?

Answer

Each GPU owns a shard of the weight matrix and computes a partial output. The runtime must gather or reduce those partial outputs so the next operation sees the same logical tensor the dense model would have produced.

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.^{[10]Reference 10NVIDIA NVLink and NVIDIA NVSwitch Supercharge Large Language Model Inferencehttps://developer.nvidia.com/blog/nvidia-nvlink-and-nvidia-nvswitch-supercharge-large-language-model-inference/} 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.

Roughly how much communication does tensor parallelism add per decode step, and why does the interconnect decide if it helps?

Answer

With the Megatron pattern, each layer needs two all-reduces in the forward pass (one for attention, one for the MLP). A deep model multiplies that by its layer count on every generated token. Fast intra-node links can lower that overhead, while slower or higher-latency paths can dominate decode and erase the benefit of adding GPUs.

Startup is only a lower bound. Each all-reduce also moves activations on the order of batch × hidden × dtype (ring factor about 2(N−1)/N for TP size N). On fast NVLink the payload term often dominates cold startup; on PCIe or cross-node links both matter:

At small activation sizes the startup floor can dominate. At larger batch or hidden sizes, or on slower links, payload and ring traffic decide whether TP=4 helps or stalls decode. Batching amortizes fixed collective startup but increases bytes moved per step.

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.^{[9]Reference 9Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.https://arxiv.org/abs/1909.08053} In serving stacks, the same idea appears as tensor_parallel_size.

When is tensor parallelism the right first sharding knob?

Answer

Use it when one request needs one model copy spread across fast-connected GPUs, especially when the model or layer tensors don't fit on one GPU and batch-only replicas can't solve the capacity problem.

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.

Why can pipeline parallelism hurt single-request latency?

Answer

Each stage depends on the previous stage. With little concurrency, later GPUs sit idle while earlier stages work, creating pipeline bubbles. The first token still has to travel through every stage.

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.^{[8]Reference 8Distributed Inference and Serving.https://docs.vllm.ai/en/latest/serving/parallelism_scaling.html} This may make a larger replica fit, but queue depth and interconnect measurements still decide TTFT and throughput.

What does tensor_parallel_size=4 and pipeline_parallel_size=2 mean on eight GPUs?

Answer

The model is split into two layer stages. Each stage is itself a four-GPU tensor-parallel group, so each stage shards large matrices across four GPUs while the full model depth is split across two stages.

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.^{[11]Reference 11Parallelism Strategies Guide.https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html}

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.^{[12]Reference 12Ring Attention with Blockwise Transformers for Near-Infinite Context.https://arxiv.org/abs/2310.01889}

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.

When does context-aware parallelism become relevant for inference?

Answer

When long prompts or attention work dominate prefill and KV memory pressure. If the main problem is model weight size, start with quantization, tensor parallelism, batching, and replicas first.

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.^{[13]Reference 13Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.https://arxiv.org/abs/1701.06538} 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.

GLM-5.2 and DeepSeek V4 Flash 0731 make this distinction concrete. GLM-5.2 routes each token to 8 of 256 experts plus one shared expert. Flash 0731 routes to 6 of 256 plus one shared expert. Serving either model still needs expert placement, routing balance, and communication measurements; expert parallelism can combine with tensor and data parallelism rather than replace them.^{[7]Reference 7GLM-5.2 Configurationhttps://huggingface.co/zai-org/GLM-5.2/blob/main/config.json[3]Reference 3DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligencehttps://arxiv.org/abs/2606.19348}

What new bottleneck does expert parallelism introduce for MoE inference?

Answer

Tokens must be routed to the GPUs that own their selected experts. If routing is imbalanced or expert communication is slow, one expert shard can bottleneck the whole decode step.

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.^{[14]Reference 14DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving.https://arxiv.org/abs/2401.09670[15]Reference 15Splitwise: Efficient Generative LLM Inference Using Phase Splitting.https://arxiv.org/abs/2311.18677} The parallelism choices below still apply, but they can be evaluated per phase.

Why is one fast 8-GPU node often a better first serving target than two weaker 4-GPU nodes?

Answer

Tensor-parallel inference communicates inside layers on every decode step. A single node with high-bandwidth links usually keeps that communication much cheaper than crossing slower network links between nodes.

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.

How do you choose between model parallelism and replicas?

Answer

Use model parallelism when a single model copy doesn't fit or one request needs multiple GPUs. Use replicas when one copy fits and traffic volume is the bottleneck. Use both when the model is large and traffic is high.

When multi-GPU serving breaks down

• 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.

What symptom suggests your multi-GPU plan is communication-bound?

Answer

Latency gets worse after sharding even though memory fits. Interconnect utilization rises, decode steps wait on collectives or activation transfers, and a smaller single-node model may answer faster.

Before moving on, trace one request through weight fit, KV-cache headroom, interconnect cost, and latency gates. Write a two-row placement checklist: "replicas because one copy fits" versus "tensor parallelism because fit requires sharding."