From Batch Size to TPU Topology: A Capacity Equation for ML Serving

When I review an ML serving capacity plan, I usually see one of two shortcuts:

replicas = peak QPS / benchmark QPS

or:

capacity = number of accelerator chips * throughput per chip

Both are useful for a quick sanity check. Neither is reliable enough for an actual plan.

The problem is the denominator. A replica does not have one fixed QPS number. Its usable throughput changes with batch size, batch timeout, input shape, model version, latency target, and accelerator topology. Replica count also changes the arrival rate seen by each batcher, which changes the batch size the server can form.

I wanted one equation that keeps those relationships visible without pretending to model every implementation detail. This is the version I use:

lambda_peak
  <= (R - F)
     * rho_slo(T)
     * C_eff(T, lambda_peak / (R - F), B_max, tau, Z)

The required replica count is the smallest integer R that satisfies this inequality.

Where:

lambda_peak is forecast peak request rate;
R is provisioned replica count;
F is replicas unavailable during the failure or rollout scenario;
T is the accelerator topology used by one replica;
rho_slo(T) is the highest utilization that still meets the latency SLO;
B_max is maximum batch size;
tau is batch timeout;
Z is the request-shape distribution;
C_eff is effective request throughput for one replica under that traffic.

Everything in the article comes back to estimating C_eff, selecting rho_slo, and comparing topology candidates.

Start With a Measured Batch Surface

For each topology T, benchmark batch service time:

S_T(b, z) = time required to execute one batch

where b is batch size and z is an input-shape bucket.

For example:

Topology	Shape bucket	Batch size	p50 batch time	p99 batch time
TPU v5e 2x2	short	1	8 ms	10 ms
TPU v5e 2x2	short	4	11 ms	14 ms
TPU v5e 2x2	short	8	16 ms	21 ms
TPU v5e 2x2	short	16	28 ms	37 ms

The corresponding saturated batch throughput is:

C_batch(T, b, z) = b / S_T(b, z)

At batch size 8 and p50 service time 16 ms:

C_batch = 8 / 0.016 = 500 requests/s

That is not yet replica capacity. It assumes the server actually executes batches of eight continuously.

The benchmark key should include model digest, runtime and compiler version, precision, topology, shape bucket, batch size, and concurrent model instances. If one of those changes, I would not reuse the number without another benchmark.

This is why I prefer a surface over one “maximum throughput” result. The plan needs to know what happens at the batch sizes production traffic will actually produce.

Effective Throughput Uses the Real Batch Distribution

Let B be the realized batch size and Z the realized shape bucket. Effective replica throughput is:

C_eff(T) = E[B] / E[S_T(B, Z)]

Expanded over observed batch and shape buckets:

             sum over b,z of P(b,z) * b
C_eff(T) = --------------------------------
             sum over b,z of P(b,z) * S_T(b,z)

This ratio matters. Capacity is completed requests divided by elapsed accelerator service time.

Do not average b / S_T(b) directly unless the weights represent time spent in each batch state. A server that executes many small, fast batches and a few large, slow batches can otherwise look more efficient than it is.

The most reliable source for P(b,z) is a production batch histogram or a representative load test. When neither exists, batch timeout gives a reasonable first approximation.

Estimating Batch Fill From Arrival Rate and Timeout

Assume one replica receives requests at rate lambda_r. When the first request enters an empty batch, the server waits up to tau for more requests, stopping early if it reaches B_max.

Under a Poisson arrival approximation:

N_tau ~ Poisson(lambda_r * tau)

B = min(B_max, 1 + N_tau)

Therefore:

P(B = k) = PoissonPMF(k - 1; lambda_r * tau)

for 1 <= k < B_max, with the remaining probability assigned to B_max.

The useful intuition is in the product:

lambda_r * tau

If a replica receives 300 requests/s and batch timeout is 3 ms:

lambda_r * tau = 300 * 0.003 = 0.9

Starting from an empty queue, that batcher should not be expected to produce batches near 16. It sees fewer than one additional arrival during the timeout on average.

This approximation is intentionally simple. When utilization is high, requests accumulate while the accelerator is busy, so the next batch may already be waiting and will be larger than this empty-queue model predicts. Bursty traffic and shape bucketing also change the distribution.

I use the Poisson result as a lower-load estimate and a consistency check. For a final plan, I use an observed batch histogram or a short queue simulation.

Replica Count Is a Fixed-Point Problem

Suppose total demand is split evenly across surviving replicas:

lambda_r = lambda_peak / (R - F)

Adding replicas lowers lambda_r. That usually improves queueing latency, but it can reduce average batch size and throughput per replica.

This is the circular dependency:

R changes lambda_r
lambda_r changes batch distribution
batch distribution changes C_eff
C_eff changes required R

That is why the core capacity equation must be solved by trying integer replica counts:

for R = F + 1, F + 2, ...:
    lambda_r = lambda_peak / (R - F)
    estimate P(B, Z) at lambda_r
    compute C_eff(T)

    if lambda_peak <= (R - F) * rho_slo(T) * C_eff(T):
        R is the first capacity-feasible replica count

Then verify latency and memory before accepting it.

This is still a small calculation. It is just more honest than dividing by a saturated throughput number once.

Choose rho_slo From the Latency Knee

The accelerator may continue processing traffic close to 100% utilization, but queueing delay becomes unstable before that point.

I define:

rho_slo(T) =
  maximum measured utilization where the required latency SLO still passes

Obtain it from a throughput-latency load test for the same model, topology, batching policy, and shape mix.

For example:

Utilization	p99 latency
50%	34 ms
60%	42 ms
70%	58 ms
75%	72 ms
80%	103 ms

If the p99 SLO is 80 ms, I might set:

rho_slo = 0.75

or slightly lower if the benchmark environment is cleaner than production.

This factor captures an important boundary: usable capacity ends where the service stops meeting its latency promise, not where the accelerator stops completing requests.

The end-to-end latency decomposition is still useful:

L_e2e =
    L_network
  + L_preprocess
  + W_batch
  + W_executor
  + S_batch
  + L_postprocess

Batch timeout directly affects W_batch. Utilization affects W_executor. Larger batches may increase S_batch while improving request throughput.

I would not add p99 values from these stages; percentiles do not add cleanly. Measure or replay end-to-end latency, then use that result to determine rho_slo.

TPU Topology Enters Through Throughput, Memory, and Chip Cost

TPU topology is not a label attached after replica count is calculated.

For a topology T:

chips(T) = product of topology dimensions

A TPU v5e 2x2 replica uses four chips. A 4x4 replica uses sixteen.

Cloud TPU supports discrete slice shapes, and topology affects model sharding, available HBM, collective communication, host boundaries, and placement. Google’s public TPU documentation describes slices as topology-shaped groups of chips connected by the inter-chip interconnect (TPU system architecture, TPU v5e configurations).

Before evaluating capacity, the topology must pass memory fit:

M_parameters(T)
+ M_runtime(T)
+ E[B] * M_per_request(Z, T)
+ M_safety
<= HBM_total(T) * eta_usable

Then benchmark S_T(b,z) on that topology.

A larger topology can reduce batch service time by spreading the model across more chips. It can also reduce throughput per chip because of communication overhead or small per-shard work. In other cases, the additional HBM permits a larger batch and materially improves throughput.

There is no safe topology multiplier. Use measured batch surfaces.

Once R(T) is known for each candidate topology, compare total chip cost:

total_chips(T) = R(T) * chips(T)

The selected topology is:

T* = argmin over T of total_chips(T)

subject to:

capacity equation passes
latency SLO passes
memory fit passes
required topology is placeable
failure scenario passes

Cost can replace chip count when accelerator generations have different prices.

Worked Example

Assume a stateless encoder service with:

peak demand:              900 requests/s
p99 latency SLO:          80 ms
max batch size:           16
batch timeout:            3 ms
unavailable replicas F:   1

The values below are synthetic.

Candidate A: TPU v5e 2x2

From the measured batch surface and expected shape mix:

Provisioned replicas	Surviving replicas	Arrival rate per survivor	Mean batch	Effective throughput per replica
3	2	450 req/s	6.8	469 req/s
4	3	300 req/s	4.9	408 req/s
5	4	225 req/s	3.6	343 req/s

Suppose load testing gives:

rho_slo(2x2) = 0.70

For R = 3:

safe capacity
  = (3 - 1) * 0.70 * 469
  = 657 requests/s

Fail.

For R = 4:

safe capacity
  = (4 - 1) * 0.70 * 408
  = 857 requests/s

Still below 900.

For R = 5:

safe capacity
  = (5 - 1) * 0.70 * 343
  = 960 requests/s

Pass.

Chip requirement:

5 replicas * 4 chips = 20 chips

Candidate B: TPU v5e 4x4

Suppose two surviving replicas each sustain effective throughput of 653 requests/s, with:

rho_slo(4x4) = 0.70

For R = 3 and F = 1:

safe capacity
  = 2 * 0.70 * 653
  = 914 requests/s

Pass.

Chip requirement:

3 replicas * 16 chips = 48 chips

Both plans pass the simplified capacity equation, but 2x2 requires fewer total chips:

20 chips vs 48 chips

The 4x4 topology may still win if it has materially better latency, availability, placement, or operational behavior. The equation narrows the decision; it does not erase product constraints.

The important observation is that the smaller topology needed more replicas but less total accelerator inventory.

For Autoregressive LLMs, Change the Work Unit

The same framework applies to LLM serving, but request QPS is usually the wrong demand unit.

Separate prefill and decode:

prefill demand = request rate * E[input tokens]
decode demand  = request rate * E[output tokens]

The benchmark surface becomes:

S_prefill(total prompt tokens, topology)

S_decode(active sequences, context distribution, topology)

The batch variable during decode is the number of active sequences or scheduled tokens, not a fixed request batch. Memory fit must include KV cache:

M_total =
    M_parameters
  + M_runtime
  + sum M_KV(context_i)
  + M_workspace
  + M_safety

The same core equation remains:

demand
  <= surviving replicas
     * SLO-safe utilization
     * effective measured throughput

Only the definition of demand and throughput changes from requests/s to tokens/s or decode steps/s.

The JAX Scaling Book’s inference chapters provide a useful theoretical treatment of the latency-throughput frontier on TPUs, including parameter bandwidth, FLOPs, KV cache, and inter-chip communication. I would use those equations to establish bounds and explain bottlenecks, then use measured runtime throughput in the capacity equation.

What Data Is Required

The calculation needs:

1. peak demand by a declared time window
2. request-shape or token distribution
3. max batch size and batch timeout
4. measured batch service-time surface
5. end-to-end latency SLO
6. rho_slo from load testing
7. topology shape and chips per replica
8. memory fit and usable HBM margin
9. replicas unavailable during failure or rollout
10. topology placement availability

If the batch histogram is available, use it directly. If not, estimate it from per-replica arrival rate and timeout, then validate with a load test.

If rho_slo is unavailable, use a conservative temporary value and label the estimate accordingly. I would rather carry visible uncertainty than hide it inside an optimistic throughput number.

Final Formula

For each candidate topology T, find the smallest integer R satisfying:

lambda_peak
  <= (R - F)
     * rho_slo(T)
     * [ E[B] / E[S_T(B, Z)] ]

with:

lambda_r = lambda_peak / (R - F)

B approximately min(
    B_max,
    1 + Poisson(lambda_r * tau)
)

or, preferably, B taken from an observed batch histogram.

Then require:

p99 end-to-end latency <= SLO
memory fit == true
topology placement == feasible

Finally choose:

T* = argmin_T [ R(T) * chips(T) ]

among candidates that pass every constraint.

That is the whole method.

It is not a perfect model of every queue, compiler decision, or traffic burst. It does force the important quantities into the same calculation: measured throughput, batch formation, topology, replica count, latency headroom, and failure capacity.

For capacity planning, that is a much stronger starting point than peak QPS divided by the best benchmark on the page.