Building a scalable batch inference system on GPUs

Here’s a practical blueprint for building a scalable batch inference system on GPUs, with the main design choices and tradeoffs.

1) Start with the workload shape

Before architecture, define:

• Model type: LLM, vision, tabular, embeddings, ranking, etc.
• Request pattern: offline batches, micro-batches, streaming queues, periodic jobs
• Latency target: seconds/minutes? or just throughput?
• SLA: completion deadlines, retries, failure tolerance
• Input size distribution: uniform vs highly variable
• Output size: fixed class labels vs large generated text
• GPU memory footprint: weights, activations, KV cache, batching overhead

This determines whether you optimize for:

• Throughput: maximize GPU utilization
• Tail latency: finish jobs quickly and predictably
• Cost efficiency: fewer GPUs, higher utilization

---

2) Reference architecture

A robust batch inference system typically looks like this:

1. Job intake

   • Receive batch request or file/job descriptor
   • Validate schema and resource requirements
   • Store metadata in DB / queue

2. Scheduler / dispatcher

   • Splits large jobs into smaller chunks
   • Assigns work to GPU workers based on:
     • available memory
     • model version
     • device utilization
     • job priority / deadline

3. GPU worker pool

   • Each worker owns 1+ GPUs
   • Loads model once and reuses it across batches
   • Performs batching, preprocessing, inference, postprocessing

4. Result aggregation

   • Merge outputs in original order
   • Persist to object storage / DB / feature store
   • Mark job complete and emit metrics

5. Observability + control plane

   • GPU utilization, memory, queue depth, batch sizes, error rates
   • Autoscaling, retry policies, job backpressure

---

3) Core scaling principles

A. Keep models hot in memory

Loading the model per request kills performance. Instead:

• preload model on worker startup
• keep workers long-lived
• rotate workers only on deploys or failures

If the model is large, consider:

• model sharding
• tensor parallelism
• quantization
• offloading non-critical parts to CPU

---

B. Batch dynamically

Static batch sizes are usually suboptimal. Use dynamic batching:

• collect requests for a short window, e.g. 5–50 ms
• build batch up to:
  • max items
  • max tokens / pixels / sequence length
  • max memory
• process batch on GPU

This improves utilization while keeping latency bounded.

For variable-length inputs:

• bucket by length
• pad minimally
• sort or pack requests to reduce waste

---

C. Chunk large jobs

For giant datasets:

• split into partitions
• process partitions independently
• write outputs incrementally
• resume from checkpoints on failure

This avoids long-running monolithic jobs and improves parallelism.

---

D. Use asynchronous pipelines

Separate stages:

• CPU preprocessing
• GPU inference
• CPU postprocessing
• storage writes

Use queues between stages so CPU and GPU can overlap. Common pattern:

• preprocessing threads prepare next batch while GPU runs current batch
• postprocessing is done asynchronously to avoid blocking GPU

---

4) GPU utilization strategies

For classic DL / CV / tabular

• Increase batch size until GPU memory or latency limit is hit
• Use mixed precision (FP16/BF16)
• Fuse kernels where possible
• Pin host memory and use async transfers
• Avoid Python overhead in tight loops

For LLM batch inference

Important considerations:

• group by prompt length / output length
• use continuous batching or in-flight batching
• manage KV cache carefully
• limit max tokens per batch, not just number of requests
• use speculative decoding if generating text at scale

Frameworks often used:

• TensorRT-LLM
• vLLM
• Triton Inference Server
• Hugging Face TGI
• DeepSpeed-Inference

---

5) Scheduling and load balancing

Scheduling policies

• FIFO: simplest, but can be unfair
• Shortest job first: better average latency, risk of starvation
• Deadline-aware scheduling: best for SLAs
• Size-aware batching: good when input lengths vary a lot

Load balancing across GPUs

Dispatch based on:

• free VRAM
• active batch size
• current queue depth
• model residency
• NUMA locality if on multi-socket hosts

If multiple models share GPUs:

• consider model pinning per GPU
• or use a centralized scheduler with admission control

---

6) Memory management

GPU memory is usually the limiting factor.

Strategies:

• pre-allocate buffers
• reuse tensors/buffers
• minimize fragmentation
• use smaller precision
• avoid copying data multiple times
• keep intermediate outputs off GPU unless needed

For LLMs:

• KV cache often dominates memory
• set strict limits on max context and max generation tokens
• evict or compact cache when possible

---

7) Autoscaling

Scale based on:

• queue depth
• estimated wait time
• GPU utilization
• p95/p99 completion time
• backlog age

Typical strategy:

• scale out when queue latency rises above threshold
• scale in only after sustained low utilization
• keep warm pool for fast bursts
• separate scaling for CPU preprocessors and GPU workers

If using Kubernetes:

• use GPU node pools
• request GPUs explicitly
• HPA/KEDA for queue-driven autoscaling
• warm pods can reduce cold start impact

---

8) Fault tolerance

Batch systems need strong recovery:

• idempotent job processing
• checkpoint completed chunks
• retry transient GPU / OOM / network errors
• dead-letter queue for failed items
• job timeouts and partial result handling

For long jobs, persist progress frequently so worker crashes don’t restart everything.

---

9) Observability

Track at least:

System metrics

• GPU utilization
• GPU memory used/free
• host CPU and RAM
• queue depth and age
• batch size distribution
• throughput (items/sec, tokens/sec)
• end-to-end latency
• error rates

Model metrics

• inference time per batch
• preprocessing time
• postprocessing time
• OOM frequency
• accuracy / quality metrics on sampled outputs

Use tracing to understand where time goes:

• enqueue → dispatch → preprocess → H2D copy → inference → D2H copy → write

---

10) Common architecture patterns

Pattern 1: Simple worker queue

Best for:

• one model
• moderate scale
• straightforward batch jobs

Components:

• queue
• GPU worker service
• storage

Pattern 2: Central scheduler + worker pool

Best for:

• many jobs
• variable input sizes
• multiple GPUs/nodes
• SLAs

Pattern 3: Serving platform for batch + online

Best for:

• shared infrastructure
• same model used for batch and real-time
• need unified observability and deployment

---

11) Practical design tips

• Prefer micro-batching over single-item inference
• Keep per-worker concurrency low enough to avoid VRAM spikes
• Benchmark on real input distributions, not just synthetic ones
• Make batch size adaptive
• Use model-specific optimizations before scaling hardware
• Separate “fast lane” and “bulk lane” for urgent vs large jobs
• Store inputs/outputs in object storage, not in memory
• Plan for partial completion and resumability

---

12) A good implementation stack

One common stack:

• Orchestration: Kubernetes
• Queue: Kafka / SQS / PubSub / Redis Streams
• Inference runtime: Triton / vLLM / custom PyTorch worker
• Storage: S3/GCS + metadata DB
• Metrics: Prometheus + Grafana
• Tracing: OpenTelemetry
• Autoscaling: KEDA / custom scaler

---

13) If you want maximum throughput

Focus on:

• large dynamic batches
• mixed precision / quantization
• minimizing CPU overhead
• overlapping pipeline stages
• model/runtime optimized for your workload
• high GPU occupancy and low fragmentation

---

14) If you want maximum reliability

Focus on:

• idempotent processing
• chunk-level checkpointing
• deterministic retries
• versioned models
• strong job metadata tracking
• graceful degradation under load

---

If you want, I can also provide one of these next:

1. a reference architecture diagram
2. a Kubernetes-based deployment design
3. a Triton/vLLM batch inference example
4. a Python worker implementation sketch
5. a scaling and autoscaling strategy for a specific model type