Plan a Multi-GPU Tensor-Parallel Mesh for LLM Serving

Pick a GPU topology and instantly see whether it is even legal — then get per-GPU weight and KV-cache shards plus the all-reduce bytes flowing across your interconnect every token.

Sharding a transformer across GPUs is not free choice. Tensor parallelism (TP) splits attention heads and the hidden dimension, so the TP degree must evenly divide both the head count and the hidden size or your matmuls will not align. Pipeline parallelism (PP) slices the layer stack, so it must divide the layer count. Data parallelism (DP) just replicates. This planner enforces those rules and quantifies the cost of the mesh you chose before you rent 32 H100s and find out it stalls on NVLink.

Total model params, e.g. 70 = 70B
= query heads if not using GQA/MQA
For KV-cache sizing
GPUs required
TP × PP × DP
Weight / GPU
params ÷ (TP × PP)
KV cache / GPU
full batch × context
Memory / GPU
weights + KV shard
All-reduce / token
across TP group, all layers
Interconnect @ batch
per decoded token, all seqs
Derived quantityValue

How the numbers are computed

The planner runs three validation checks and four sizing formulas entirely in your browser. Validation: TP must divide the query-head count, the KV-head count, and the hidden size; PP must divide the layer count; TP × PP × DP is the total device count. If any divisibility fails the mesh is flagged invalid — those matmuls cannot be evenly partitioned. A warning is raised when a shard is badly imbalanced (for example TP not dividing head_dim cleanly) even if technically legal.

weight_per_gpu = (params × dtype_bytes) / (TP × PP) kv_bytes_token = 2 × layers × kv_heads × head_dim × dtype_bytes (head_dim = hidden / q_heads) kv_per_gpu = kv_bytes_token × seq_len × batch / (TP × PP) allreduce_token = 2 × (TP − 1) / TP × hidden × dtype_bytes × (layers/PP) × 2

Weight memory divides by TP × PP because tensor parallelism shards each layer's matrices while pipeline parallelism assigns disjoint layer groups to different stages — both reduce a single GPU's parameter footprint. The KV cache is likewise sharded across the same TP × PP group, and it scales with the KV-head count (GQA and MQA shrink it dramatically versus full multi-head attention). The all-reduce term is the classic ring cost: each transformer layer performs two all-reduces (after the attention projection and after the MLP), and a ring all-reduce moves 2(TP−1)/TP of the activation tensor per participant. Multiplying by layers-per-stage gives the per-token cross-GPU traffic that your NVLink or InfiniBand fabric must sustain — the number that decides whether a topology is throughput-bound. Use it to compare an 8×TP single node against a 4×TP × 2×PP split before committing hardware.

Related Tools