Distributed Inference Ledger: When the Model Doesn't Fit, Where Does State Go

Inference differs from training: training only needs throughput; inference needs latency and state management. When the model doesn't fit one card, how you shard it, where state lives, and how you scale out — every step has a memory and bandwidth ledger. This article works through it. Method framework references the open-source book Understanding AI Infra (Apache-2.0); formulas and examples are derived by this site and are reproducible.

1. Memory Ledger: Doesn't Fit → Tensor Parallel

Tensor parallelism (TP) splits weights along the hidden dimension; N cards each hold 1/N:

$$Mem\_per\_card = \frac{W}{TP} + \frac{KV}{TP} + activation\_headroom$$

70B FP16 (140 GB weights) + 32K context KV (≈5 GB): at TP4 each card ≈ 36 GB — from "doesn't fit" to "runnable". That is the quantitative reason 70B can run on 48G cards.

2. Comm Ledger: TP's Hidden Cost

TP requires All-Reduce before and after every op (two syncs):

$$Comm\_per\_token \approx 2 \times \frac{TP-1}{TP} \times H_{hidden} \times bytes \times layers(\approx)$$

Per-token TP traffic scales with hidden dim × layers, independent of batch size. Therefore:

• TP is for not-enough-memory scenarios, not a "faster" tool;
• at small batch, TP comm ratio is high; at large batch, data parallel (DP) wins;
• Takeaway: ask "is memory enough?" first; only if not, use TP. If memory is fine, prefer DP.

Scheme	Mem/card	Per-token comm	Fits
DP (full replica)	W+KV full	none (request routing only)	enough memory, high concurrency
TP2	half	1×H×bytes×layers	memory-critical
TP4	1/4	1.5×H×bytes×layers	big model, insufficient memory
TP8	1/8	1.75×H×bytes×layers	very large model

3. State Ledger: Where KV Lives

KV cache is "session state"; the core problem of distributed inference is state placement:

Placement	Pros	Cost
Same card as weights (TP cards)	no extra transfer	occupies VRAM, limits concurrency
Dedicated KV node / offload	higher weight-layer utilization	fetch KV cross-node per step
Prefill/decode separation (PD)	each stage optimized independently	state migration (saved on prefix-cache hits)

The PD ledger: prefill is compute-bound (one pass over a long input), decode is bandwidth-bound (reads all weights per token) — their optimal configurations differ. Splitting prefill and decode onto separate machine pools, each tuned to its own bottleneck, is the standard way to scale inference (see Inference Engine Tuning and vLLM production deployments).

4. Scale-Out & Recovery Ledger

• Stateless part (weights): add replicas anytime — N more DP replicas → throughput ≈ ×N, provided requests can be spread;
• Stateful part (KV): hard to scale — in-flight sessions can't migrate at will; prefix-cache hits allow "continuation";
• Recovery: a failed card drops sessions on it; stateless replicas backfill automatically; KV rebuilt from snapshots.

5. Reproducible Python

def tp_account(n_params_b, kv_gb, tp, bytes_per_param=2, h_dim=8192, layers=80):
    """TP inference mem/card and per-token comm (MB)"""
    mem = n_params_b * 1e9 * bytes_per_param / tp + kv_gb / tp
    comm_mb = 2 * (tp - 1) / tp * h_dim * layers * 2 / 1e6
    return mem, comm_mb

tp_account(n_params_b=70, kv_gb=5.2, tp=4)
# (36.3 GB/card, 2.4 MB/token)

6. Checklist

1. Do the memory ledger first to choose DP/TP: memory fine → DP first; not enough → TP;
2. Start TP at TP2 and measure comm ratio — don't jump to TP8 (comm ledger grows with TP);
3. Define KV placement explicitly: same-card / offload / PD separation based on concurrency & latency targets;
4. Promise scale-out only for the stateless part (DP replicas); estimate stateful gains by prefix-hit rate;
5. Drill a single-card-failure recovery once, and confirm session-drop & backfill behavior matches expectations.

---

Derived and organized by this site; method framework references the open-source book Understanding AI Infra (Apache-2.0). Version numbers and measurements are subject to on-site conditions.