Step 1: Batch size sweep
We fix a small transformer encoder forward pass (typical production inference shape) and sweep batch sizes ∈ {1, 4, 16, 64}. At each size, we measure:
- Throughput (samples/s) — the main deliverable.
- Peak VRAM — what you'd reserve in Kubernetes resource requests.
- Latency per sample — what users feel.
Three things to watch as batch grows:
- Throughput climbs steeply at first (overhead amortization), then flattens.
- Latency per batch climbs; latency per sample falls.
- VRAM grows roughly linearly (activations scale with batch).
The model + fixed constants we'll sweep against
A tiny transformer encoder — small enough that batch=1 is launch-overhead-bound, which is exactly the regime where batching wins the most. Constants are module-level so they stay stable across every cell in this lab; if you fork this lab to a larger model, just bump D_MODEL/N_LAYERS.
import torch, torch.nn as nn, time, warnings
warnings.filterwarnings('ignore', message='.*enable_nested_tensor.*')
device = 'cuda'
D_MODEL = 256 # small enough that batch=1 is launch-overhead-bound
SEQ_LEN = 64
N_LAYERS = 2
VOCAB = 16000
ITERS = 100
WARMUP = 10
class SmallTransformer(nn.Module):
"""Tiny transformer encoder sized for clean batch-scaling measurements on consumer GPUs."""
def __init__(self, d=D_MODEL, nlayers=N_LAYERS, seq_len=SEQ_LEN, vocab=VOCAB):
super().__init__()
self.embed = nn.Embedding(vocab, d)
layer = nn.TransformerEncoderLayer(
d_model=d, nhead=4, dim_feedforward=4*d,
dropout=0.0, batch_first=True, norm_first=True,
)
self.enc = nn.TransformerEncoder(layer, num_layers=nlayers)
self.head = nn.Linear(d, vocab)
def forward(self, x):
return self.head(self.enc(self.embed(x)))
The bench harness
bench(model, batch, dtype) returns (throughput, peak_vram, latency). Two details matter:
- Warmup before measurement — first iterations pay JIT + workspace allocation costs.
torch.cuda.synchronize()beforeperf_counter()— the CPU-side timer is wall-clock, but CUDA calls are async. Without the sync, you're timing submission, not execution.
def bench(model, batch, dtype=torch.float32):
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
model = model.to(device=device, dtype=dtype)
x = torch.randint(0, VOCAB, (batch, SEQ_LEN), device=device)
with torch.no_grad():
for _ in range(WARMUP):
y = model(x)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(ITERS):
y = model(x)
torch.cuda.synchronize()
elapsed = time.perf_counter() - t0
peak = torch.cuda.max_memory_allocated() // (1024 * 1024)
samples_per_s = ITERS * batch / elapsed
latency_ms = 1000 * elapsed / ITERS
return samples_per_s, peak, latency_ms
Batch sweep — find the knee
Sweep batch sizes 1 → 64 and record throughput. The knee is the last batch size where doubling still gave at least a 10% throughput improvement; beyond the knee you're paying memory without a speed return. This is the single most important number for serving capacity planning.
model = SmallTransformer()
batch_sweep = []
for bs in [1, 4, 16, 64]:
tp, vram, lat = bench(model, bs, dtype=torch.float32)
batch_sweep.append({
'batch_size': bs,
'throughput_samples_per_s': tp,
'peak_vram_mib': vram,
'latency_ms': lat,
})
print(f" batch={bs:3d}: {tp:8.1f} samples/s | peak VRAM {vram:5d} MiB | latency {lat:5.2f} ms/batch")
# The knee: the batch where throughput stops growing meaningfully (<10% gain from previous)
knee_batch_size = batch_sweep[0]['batch_size']
for prev, cur in zip(batch_sweep[:-1], batch_sweep[1:]):
gain = (cur['throughput_samples_per_s'] - prev['throughput_samples_per_s']) / prev['throughput_samples_per_s']
if gain >= 0.10:
knee_batch_size = cur['batch_size']
print(f"\nKnee found at batch={knee_batch_size} (last size where doubling gave >=10% gain)")
