Serving embeddings at scale with ONNX Runtime

Last quarter, my team took over the real-time feature generation pipeline for our order-book anomaly detection models. The core of the pipeline relies on generating dense embeddings from sequence data (order-book deltas, trade prints, and quote updates) to feed downstream trade-execution models.

When we inherited the service, it was running raw PyTorch inside a FastAPI application deployed across 12 GPU-enabled Kubernetes nodes. The service was choking: at 2,500 requests per second (RPS), our p99 inference latency soared to 340ms, causing cascading timeouts in our execution routers. Worse, the GPU utilization hovered at a pathetic 14% because of Python GIL contention and unoptimized tensor operations.

We rebuilt the service using ONNX Runtime (ORT) coupled with dynamic batching, tokenization optimizations, and CUDA execution provider tuning. We eliminated PyTorch from the production path, dropped our p99 latency to 18ms at 8,500 RPS, and decommissioned 8 of our 12 GPU nodes.

Here is the exact architecture, the technical dead-ends we hit along the way, and the code you need to replicate these results.

The Bottleneck: Why PyTorch in Production Failed Us

Our baseline setup used Hugging Face’s transformers library directly in PyTorch (torch.no_grad()) to compute 384-dimensional embeddings from mini-batches of market data sequences.

When we profiled the Python processes using py-spy and NVIDIA Nsight Systems, three bottlenecks dominated the trace:

1. GIL and Web Framework Overhead: FastAPI workers spent almost 35% of their execution time marshaling JSON payloads, converting Python lists to PyTorch tensors, and waiting on the Python Global Interpreter Lock (GIL).
2. Suboptimal Graph Execution: PyTorch’s dynamic computational graph introduced non-negligible interpreter overhead per forward pass. Kernels were being launched individually without aggressive layer fusion (e.g., combining LayerNorm and Attention operations).
3. Memory Allocation Churn: Repeatedly allocating GPU memory buffers for dynamic batch sizes caused memory fragmentation and synchronous host-to-device copy delays.

Our first impulse was to scale horizontally by adding more pods. That drove our AWS bill through the roof without fixing the long-tail latency p99 spikes caused by PyTorch’s CUDA stream synchronized calls. We needed a static, highly optimized graph representation and a zero-Python inference loop.

The Target Architecture

We shifted to a two-phase architecture: an offline export pipeline that bakes static optimizations into an ONNX graph, and an asynchronous C++ runtime wrapper invoked via a thin Python interface (using onnxruntime-gpu).

flowchart LR
 client["Client Request"] -->|"gRPC Payload"| gateway["API Gateway"]
 gateway -->|"Raw Strings"| tok["Fast Tokenizer"]
 tok -->|"Int32 Tensors"| batcher["Dynamic Batcher"]
 batcher -->|"Batched Tensors"| ort["ONNX Runtime"]
 ort -->|"CUDA Execution"| gpu["NVIDIA Tensor Core"]
 gpu -->|"FP16 Embeddings"| client

By decoupling tokenization from tensor execution and leveraging ONNX Runtime’s custom CUDA allocations, we created a pipeline where incoming sequences stay entirely in continuous memory blocks from request receipt to embedding extraction.

Step 1: Exporting and Optimizing the Graph

Exporting a model from PyTorch to ONNX sounds trivial using torch.onnx.export, but if you rely on default settings, you end up with an unoptimized graph that performance-regresses compared to native PyTorch.

Here is the script we built to export a SentenceTransformers-style model (all-MiniLM-L6-v2) to ONNX with dynamic axes and custom ONNX opset definitions.

Graph Optimization and FP16 Quantization

Exporting is only half the battle. Raw ONNX graphs contain redundant node operations. We used the onnxruntime.transformers optimizer to fuse Multi-Head Attention blocks, GELU activations, and LayerNorm layers into optimized single CUDA kernels.

When running optimize_onnx.py, inspect the node count reduction:

Going from 312 nodes to 94 reduced our CUDA kernel invocation overhead by over 60%.

Step 2: High-Performance Production Engine

A fast graph is useless if your wrapper code introduces standard Python latency overhead. Below is our complete production-grade embedding server pipeline implementing tokenization, pinned memory allocations, mean pooling, and ONNX Runtime execution provider configuration.

Real Dead-Ends We Encountered

We did not hit 8,500 RPS on our first try. Here are the three main failures that cost us days of debugging:

1. The IOBinding Trap with Variable Sequence Lengths

ONNX Runtime supports an advanced API called IOBinding, which allows pre-allocating GPU tensors to avoid memory copy overhead between Host (CPU) and Device (GPU).

We spent three days building an IOBinding queue. The result? Performance regressed by 45%.

Because market sequence payloads fluctuated dramatically in size (from 12 tokens to 128 tokens per batch), pre-allocated static GPU memory blocks had to be destroyed and re-allocated on almost every single request when shapes mismatched. IOBinding shines only when your batch dimensions and sequence lengths are 100% static. If your inputs have dynamic dimensions, standard host-to-device streaming via ONNX Runtime’s dynamic allocation pool is faster.

2. FP16 Underflow in L2 Normalization

When we converted our model directly to FP16, we started noticing occasional NaN embeddings in our output logs.

During L2 normalization, we compute the sum of squares ($\sum x_i^2$). In FP16, the maximum representable value is $65,504$. For a 384-dimensional vector, if multiple values exceeded $14.0$, intermediate square aggregations overflowed to $\infty$, yielding NaN after division.

The Fix: Keep model output vectors in FP16, but cast them to FP32 during the mean-pooling and L2-normalization compute steps in NumPy/C++.

3. CPU Core Over-Subscription

Initially, we ran 8 Gunicorn worker processes per container, assuming it would maximize GPU pipe throughput. Instead, throughput collapsed under high load. The issue was intra_op_num_threads. Each ONNX Runtime session spawned internal worker threads that fought over CPU cycles with Gunicorn workers.

Setting intra_op_num_threads = 2, inter_op_num_threads = 2, and deploying 2 workers per pod completely eliminated thread thrashing and dropped CPU context switches by 82%.

Benchmark Results: PyTorch vs. ONNX Runtime

We benchmarked our production setup using locust driving traffic against an AWS g5.xlarge instance (1x NVIDIA A10G GPU, 4 vCPUs, 16GB RAM). We tested across varying batch sizes with a fixed sequence length of 64 tokens.

Latency Profile under 5,000 RPS Load

Throughput vs. Latency Curve

By leveraging fused graph operations and low-overhead thread scheduling, ONNX Runtime sustained 8,500 RPS within our strict < 20ms p99 SLO.

Key Takeaways

1. Export is not Optimization: Merely running torch.onnx.export leaves vast amounts of performance on the table. You must perform explicit node fusion (e.g., combining attention heads) using ONNX model optimizers.
2. Watch your Math in FP16: Convert intermediate reductions (norm calculations, accumulators) back to FP32 to avoid numerical overflow and downstream NaN vectors.
3. Keep Dynamic Axes Dynamic: Don’t force IOBinding unless your tensor shapes are strictly fixed. The dynamic CUDA memory allocator inside ONNX Runtime handles dynamic sequence lengths efficiently without manual intervention.
4. Bench Your Framework Wrapper: A significant chunk of latency in production Python ML services lives outside the forward pass. Replacing native Transformers tokenizers with the Rust-backed tokenizers library and streamlining NumPy vector operations is critical to unlocking raw GPU hardware throughput.