Multi-GPU serving

1 · Two orthogonal axes advanced

Once a model is in production you scale it along two independent axes, and conflating them is the most common multi-GPU mistake. Replicas add throughput and availability: each replica is a complete copy of the model, and a load balancer fans requests across them. Sharding (tensor / pipeline parallelism) splits one model across several GPUs because it does not fit on one — it buys capacity to run the model at all, not throughput. You size them separately: sharding first (can one copy even load?), then replicas (how many copies for the traffic?).

A production deployment is the product of the two: total GPUs = replicas × GPUs-per-replica. Eight GPUs might be 8 replicas of a model that fits one card, or 2 replicas of a model sharded across 4 cards — very different throughput and cost profiles for the same hardware bill.

Read it left to right: a request hits the router / load balancer, which picks one of N replicas. Inside each replica the model is sharded across G GPUs with tensor parallelism — the horizontal axis (N) is throughput/availability, the depth inside a box (G) is the fit-the-model axis. The two multiply.

2 · Tensor parallelism at inference advanced

Tensor parallelism splits each weight matrix across G GPUs: every GPU holds a slice of each layer and computes a partial result on the full activation. Because a layer's output is a sum over those slices, the GPUs must all-reduce (sum-and-broadcast) their partials twice per transformer layer — once after attention, once after the MLP. That collective runs over NVLink and adds latency to every token: TP raises capacity but never comes free. It is why TP is kept within a single node (fast NVLink) and pipeline parallelism spans nodes (slower links, tolerates it).

gpus_needed.pydef gpus_needed(model_gb, gpu_gb, kv_overhead=1.4, usable=0.90):
    """GPUs to hold weights + KV-cache/activations.
    kv_overhead: multiply weights to reserve room for KV-cache + activations.
    usable: fraction of each GPU you dare use (rest = fragmentation/headroom).
    Tensor parallelism needs a power-of-two shard count in practice."""
    need_gb = model_gb * kv_overhead
    per_gpu = gpu_gb * usable
    raw = -(-need_gb // per_gpu)            # ceil division
    g = 1
    while g < raw:                          # round up to a power of two
        g *= 2
    return {"need_gb": round(need_gb, 1), "per_gpu_gb": round(per_gpu, 1),
            "raw_gpus": int(raw), "tensor_parallel_size": g}

# Llama-3.1-70B in fp16 is ~140 GB of weights; an A100/H100 is 80 GB.
print(gpus_needed(model_gb=140, gpu_gb=80))
# A 4-bit (AWQ) 70B is ~40 GB and fits two cards.
print(gpus_needed(model_gb=40, gpu_gb=80))

{'need_gb': 196.0, 'per_gpu_gb': 72.0, 'raw_gpus': 3, 'tensor_parallel_size': 4}
{'need_gb': 56.0, 'per_gpu_gb': 72.0, 'raw_gpus': 1, 'tensor_parallel_size': 1}

3 · KV-cache is your capacity ceiling expert

Weights decide whether the model loads; the KV-cache decides how many requests you can serve at once. Every in-flight request holds a per-token key/value tensor for every layer for its whole context. Once the weights are resident, all remaining GPU memory is a KV-cache budget — and that budget, divided by per-request KV size, is the hard cap on concurrency. Run past it and vLLM preempts/queues requests; latency spikes. This is the number that actually limits a serving box, not FLOPs.

kv_capacity.pydef per_request_kv_gb(layers, kv_heads, head_dim, seq_len, bytes_per=2):
    """KV bytes = 2 (K and V) * layers * kv_heads * head_dim * seq_len * dtype_bytes.
    kv_heads is the GROUPED count (GQA/MQA) — far smaller than attention heads."""
    kv_bytes = 2 * layers * kv_heads * head_dim * seq_len * bytes_per
    return kv_bytes / 1e9

def max_concurrent_requests(kv_cache_budget_gb, per_request_kv_gb):
    """Concurrency ceiling = KV budget / per-request KV footprint."""
    return int(kv_cache_budget_gb // per_request_kv_gb)

# Llama-3-70B shape: 80 layers, 8 GQA KV-heads, head_dim 128, 8k context, fp16.
one_req = per_request_kv_gb(layers=80, kv_heads=8, head_dim=128, seq_len=8192)
print(f"per request @8k ctx: {one_req:.2f} GB")

# After 140 GB of weights across 4x80GB cards, ~116 GB is left for KV.
print("max concurrent:", max_concurrent_requests(kv_cache_budget_gb=116, per_request_kv_gb=one_req))
# Halve the context and you nearly double concurrency:
half = per_request_kv_gb(layers=80, kv_heads=8, head_dim=128, seq_len=4096)
print("max concurrent @4k:", max_concurrent_requests(116, half))

per request @8k ctx: 2.68 GB
max concurrent: 43
max concurrent @4k: 86

4 · vLLM / TGI multi-GPU flags expert

You configure parallelism; you never implement it. In vLLM the shard count is --tensor-parallel-size (a.k.a. tensor_parallel_size in the LLM(...) Python API), with --pipeline-parallel-size for spanning nodes. In TGI (Text Generation Inference) the equivalent is --num-shard. Both expose --gpu-memory-utilization (vLLM) / --cuda-memory-fraction (TGI) to tune the weights-vs-KV split from section 3. The command below is real but needs the hardware.

serve_tp4.sh# ▶ needs multiple GPUs (>=4) on one node with NVLink + `pip install vllm`.
# NOT run by this lesson's offline labs — it talks to real hardware.

# Serve a 70B model sharded 4 ways (tensor parallel), OpenAI-compatible on :8000
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 \          # split each layer across 4 GPUs (all-reduce/layer)
  --gpu-memory-utilization 0.90 \     # leave the rest as KV-cache budget (section 3)
  --max-num-seqs 128 \                # continuous-batch width; bounded by KV, not this
  --host 0.0.0.0 --port 8000

# TGI equivalent shard flag would be:  text-generation-launcher --num-shard 4

# --- client: identical to any OpenAI call, one replica behind the port ---
# from openai import OpenAI
# client = OpenAI(base_url="http://localhost:8000/v1", api_key="none")
# r = client.chat.completions.create(
#     model="meta-llama/Llama-3.1-70B-Instruct",
#     messages=[{"role": "user", "content": "Hello from 4 sharded GPUs."}])
# print(r.choices[0].message.content)

Note what the flags do not do: --tensor-parallel-size 4 makes the 70B model fit and share its layers across 4 cards; it does not multiply throughput. To serve more traffic you run this same command on more nodes and put a load balancer in front — that is the replicas axis, orchestrated outside vLLM (Kubernetes, a router).

5 · Autoscaling & load balancing professional

Replicas are stateless copies, so a plain load balancer fans requests across them — but for LLM serving, least-outstanding-requests (or queue-aware) routing beats round-robin, because request cost varies wildly with output length and a naive round-robin piles long generations onto one replica. Autoscale on the right signal: CPU is useless (GPUs are busy while CPU idles); scale on queue depth, in-flight requests, or p95 latency / time-to-first-token. Each new replica is a full model reload (tens of seconds to minutes), so scale ahead of demand and keep warm headroom — you cannot cold-start a 70B replica inside a latency SLO.

topology_cost.pydef plan_cost(name, replicas, gpus_per_replica, rps_per_replica,
              gpu_hourly, target_rps):
    """Compare two topologies for the SAME model+hardware budget.
    Sharding (high gpus_per_replica) fits big models but each replica is pricey
    and TP overhead caps its rps; replicas multiply throughput linearly."""
    total_gpus = replicas * gpus_per_replica
    capacity_rps = replicas * rps_per_replica
    cost_hr = total_gpus * gpu_hourly
    tokens_hr = capacity_rps * 3600 * 300          # assume ~300 tok/response
    cost_per_1k = cost_hr / (tokens_hr / 1000) if tokens_hr else float("inf")
    return {"plan": name, "gpus": total_gpus, "cap_rps": capacity_rps,
            "meets_target": capacity_rps >= target_rps,
            "cost/hr": round(cost_hr, 2), "cost/1k_tok": round(cost_per_1k, 5)}

TARGET, PRICE = 40, 4.0                              # rps needed, $/GPU-hr
# Plan A: model fits 1 GPU -> 8 lean replicas (TP overhead = none, high rps each)
print(plan_cost("A: 8x TP1", replicas=8, gpus_per_replica=1,
                rps_per_replica=6, gpu_hourly=PRICE, target_rps=TARGET))
# Plan B: same 8 GPUs but model sharded 4-way -> 2 replicas, TP tax lowers rps each
print(plan_cost("B: 2x TP4", replicas=2, gpus_per_replica=4,
                rps_per_replica=9, gpu_hourly=PRICE, target_rps=TARGET))

{'plan': 'A: 8x TP1', 'gpus': 8, 'cap_rps': 48, 'meets_target': True, 'cost/hr': 32.0, 'cost/1k_tok': 0.00062}
{'plan': 'B: 2x TP4', 'gpus': 8, 'cap_rps': 18, 'meets_target': False, 'cost/hr': 32.0, 'cost/1k_tok': 0.00165}

Same 8 GPUs, same hourly bill — but if the model fits one card, 8 lean replicas serve ~2.7× the traffic at ~2.7× lower cost-per-token than 2 sharded replicas. Sharding is what you do when Plan A is impossible because the model won't fit — not a topology you pick for speed.

6 · Tech-lead — choosing the topology tech-lead

A lead owns the whole sizing decision and writes it down: (1) can one replica load? — use gpus_needed to set the minimum shard count, preferring 4-bit weights and a single node so TP stays on NVLink; (2) what concurrency does one replica sustain? — bound it with the kv_capacity model at the real p95 context length, not the max; (3) how many replicas for the traffic + an availability margin? — from the load test, plus one for headroom; (4) autoscale on queue depth with warm capacity because replicas cold-start slowly; (5) cost per 1k tokens from topology_cost is the number you defend to the business. The senior instinct: shard only as much as you must, replicate for everything else.

def axis(symptom):
    if symptom == "model won't load on one GPU":
        return "tensor parallel (split the model, tp_size > 1)"
    if symptom == "load too high, model fits one GPU":
        return "more replicas (horizontal scale)"
    if symptom == "big model AND high load":
        return "both: TP within a replica, then replicate the replica"
    return "measure first"

for s in ["model won't load on one GPU", "load too high, model fits one GPU"]:
    print(s, "->", axis(s))

def tp_weight_gib(params_b, tp, bytes_per=2):
    return params_b * 1e9 * bytes_per / tp / (1024 ** 3)

def tp_comm_gib_per_tok(n_layers, hidden, tp, bytes_per=2):
    # one all-reduce per layer over the hidden vector, ~2*(tp-1)/tp volume
    per = 2 * (tp - 1) / tp * hidden * bytes_per
    return n_layers * per / (1024 ** 3)

for tp in (1, 2, 4, 8):
    w = tp_weight_gib(70, tp)
    c = tp_comm_gib_per_tok(80, 8192, tp)
    print(f"tp={tp}: weights/GPU={w:6.1f} GiB, comms={c*1e6:6.2f} KiB/token")
# weights fall as 1/tp; comms rises with tp -- fast interconnect (NVLink) required

def kv_per_req_gib(n_layers=80, n_kv_heads=8, head_dim=128, seq=4096, bytes_per=2):
    return 2 * n_layers * n_kv_heads * head_dim * seq * bytes_per / (1024 ** 3)

def concurrency(card_gib=80, weight_gib=32.6, headroom=0.10, seq=4096):
    budget = card_gib * (1 - headroom) - weight_gib
    per = kv_per_req_gib(seq=seq)
    return max(0, int(budget // per)), per

for seq in (2048, 4096, 8192):
    n, per = concurrency(seq=seq)
    print(f"ctx={seq:>5}: KV/req={per:.2f} GiB -> ~{n} concurrent requests")
# doubling context halves concurrency -- the KV cache IS the capacity ceiling

def validate(tp, max_model_len, gpu_mem_util, max_num_seqs,
             params_b=70, card_gib=80, bytes_per=2):
    weight = params_b * 1e9 * bytes_per / tp / (1024 ** 3)
    budget = card_gib * gpu_mem_util - weight
    kv_per = 2 * 80 * 8 * 128 * max_model_len * bytes_per / (1024 ** 3)
    need = kv_per * max_num_seqs
    ok = need <= budget
    return (f"tp={tp} len={max_model_len} util={gpu_mem_util} seqs={max_num_seqs}: "
            f"KV need={need:.1f} / budget={budget:.1f} GiB -> "
            f"{'OK' if ok else 'WILL OOM -- lower max_num_seqs or max_model_len'}")

print(validate(tp=4, max_model_len=8192, gpu_mem_util=0.90, max_num_seqs=64))
print(validate(tp=4, max_model_len=8192, gpu_mem_util=0.90, max_num_seqs=8))

# needs GPU -- the real launch (documented vLLM CLI):
# vllm serve meta-llama/Llama-3.1-70B --tensor-parallel-size 4 #   --max-model-len 8192 --gpu-memory-utilization 0.90 --max-num-seqs 8

def autoscale(replicas, kv_util, queue, up=0.80, down=0.40, max_r=20):
    if kv_util > up or queue > 5:
        return min(max_r, replicas + 1), "scale UP (KV/queue high)"
    if kv_util < down and queue == 0 and replicas > 1:
        return replicas - 1, "scale DOWN (idle)"
    return replicas, "hold"

r = 2
for kv, q in [(0.85, 2), (0.9, 8), (0.5, 0), (0.3, 0)]:
    r, why = autoscale(r, kv, q)
    print(f"kv={kv} queue={q} -> replicas={r} ({why})")

import math

def kv_per_req(seq=8192):
    return 2*80*8*128*seq*2 / (1024**3)

def topology(users, params_b=70, card_gib=80, tp=4, gpu_hr=2.0, seq=8192):
    weight = params_b*1e9*2/tp/(1024**3)
    budget = card_gib*0.9 - weight
    per_replica_users = max(1, int(budget // kv_per_req(seq)))
    replicas = math.ceil(users / per_replica_users)
    gpus = replicas * tp
    monthly = gpus * gpu_hr * 24 * 30
    return (f"tp={tp} x {replicas} replicas = {gpus} GPUs; "
            f"{per_replica_users} users/replica; ~${monthly:,.0f}/mo")

print(topology(200))
# p95 latency must be MEASURED under load (locust/vegeta + real weights) -- needs GPU