I Rebuilt an AI Gateway in Rust. It Became 41× Faster.
Reimplementing LiteLLM with Rust with h5i-app, optimizing streaming and database operations, and formally verifying core logic with Lean 4.
As applications increasingly rely on multiple LLM providers, developers need a way to manage API keys, route requests, enforce rate limits, track usage, and control costs through a unified interface. This is why AI gateways, services that unify various LLM inference APIs, are becoming an essential part of modern AI infrastructure.
One of the most popoular open-source AI gateways is LiteLLM, a Python-based AI gateway that supports a wide range of LLM providers.
Last weekend, one simple question came to my mind: “What if we rebuilt an AI gateway in Rust? How much performance could we gain?” .
To find out, I built h5i-ai-gateway, an open-source Rust implementation of an AI gateway inspired by LiteLLM, using the Rust web framework, h5i-app which is a Axum-based framework backed with Lean 4 formal verification.
The results was quite promising:
- 41.4× higher throughput for standard chat completions
- 40.5× higher throughput for streaming responses
- 30.1× higher throughput for embeddings
- 65× lower memory usage, from 4.4 GB to just 68 MB
- More than 100× improvement in p99 latency, from 346 ms to 3.08 ms
h5i-app framework also alloed me to formally verify important properties like the correctness of authentication with Lean 4.
What Is an AI Gateway?
An AI gateway sits between your application and LLM providers such as OpenAI, Anthropic, and Google. Instead of integrating each provider independently, applications send requests through a single gateway.
Application
|
v
AI Gateway
|
+--------+--------+
| | |
v v v
OpenAI Anthropic GeminiAn AI gateway typically handles several responsibilities:
- Authentication: Validate API keys and control access to models.
- Routing: Forward requests to the appropriate LLM provider.
- Rate limiting: Prevent excessive requests from individual users or teams.
- Budget management: Track token usage and enforce spending limits.
- Streaming: Relay responses using Server-Sent Events (SSE).
While his additional infrastructure is worthwhile for many applications because it centralizes important operational concerns, the gateway itself can become a significant source of overhead as AI agents begin generating hundreds or thousands of LLM requests,
(Re-) Building the Gateway in Rust
I re-implemented the LiteLLM gateway in Rust, using h5i-app, a Rust web framework built on top of Axum.
One of the central design principles of h5i-app is to separate application logic from HTTP infrastructure. Instead of embedding authorization, budgeting, and routing decisions directly into HTTP handlers, the gateway implements these operations in a separate Rust kernel. This architecture also allows the core logic to be translated into Lean 4 for formal verification.
The overall structure looks like this:
HTTP Request
|
v
Gateway Server
|
v
Rust Kernel
|
+------+------+
| |
Admit Settle
| ^
v |
Provider API -----+
|
v
HTTP ResponseEach request goes through two important stages.
1. Admit: Validate and reserve
Before forwarding a request, the gateway checks:
- Whether the API key is valid
- Whether the requested model is allowed
- Whether the request satisfies rate limits
- Whether sufficient budget is available
If the request is accepted, the gateway reserves the estimated cost.
2. Settle: Finalize usage
After the provider returns a response, the gateway calculates the actual cost using token usage information and updates the corresponding budget.
Get Hideaki Takahashi’s stories in your inbox
Join Medium for free to get updates from this writer.
This separation makes it easier to reason about the correctness of accounting operations, especially when multiple requests execute concurrently. The gateway also supports chat completions, embeddings, streaming responses, and routing across multiple providers.
Benchmark: Rust vs. Python LiteLLM
To make the performance comparison meaningful, I wanted to measure the gateway itself rather than the speed of the underlying LLM. Specifically, I compared the original Python implementation of LiteLLM and my Rust re-implementation, using 8 CPUs and 64 concurrent connections.
Throughput
The Rust implementation achieved more than 40 times the throughput for both standard chat completions and streaming.
However, the PostgreSQL-backed workload showed a much smaller improvement, showing that a faster language does not automatically eliminate database bottlenecks.
Memory Usage
The difference in memory consumption was equally significant.
The Rust implementation used approximately 65 times less memory in this benchmark configuration. This can become particularly important when deploying gateways across multiple replicas or operating resource-constrained infrastructure.
Latency
For the chat completion benchmark, the latency results were:
The improvement in tail latency was particularly impressive. At p99, latency decreased from 346 ms to just 3.08 ms. Of course, real LLM requests often spend much longer waiting for inference. To explore this, I also benchmarked the gateways with an artificial 50 ms upstream delay.
At 256 concurrent connections:
- Python LiteLLM: 657 req/s
- Rust gateway: 4,887 req/s
The Rust implementation still achieved approximately 7.4× higher throughput. These results suggest that gateway overhead can remain significant even when upstream requests involve network delays.
What Actually Made It Faster?
The performance improvement was not simply the result of replacing Python with Rust, and several optimizations made a substantial difference.
1. Optimizing TCP for Streaming Responses
LLM streaming typically uses Server-Sent Events (SSE), where the server continuously sends small chunks of data to the client. Unlike a conventional HTTP response, which might return a single large payload, an SSE connection can involve many small writes. These writes can interact with TCP buffering and delayed acknowledgments, introducing unnecessary latency.
One important optimization was enabling TCP_NODELAY on accepted TCP sockets. This disables Nagle’s algorithm, allowing small TCP segments to be transmitted without waiting for additional data to accumulate.
In the streaming benchmark with 64 concurrent connections, enabling this option improved throughput from approximately: 1,550 req/s → 11,600 req/s (roughly a 7.5x improvement).
2. Batching PostgreSQL Operations with Group Commit
AI gateways often need to maintain accurate accounting information. For example, a request may reserve part of a user’s remaining budget before execution, then settle the final amount after the response is received.
In a PostgreSQL-backed implementation, these operations typically require transactions to preserve consistency. However, issuing an independent transaction for every individual update can become expensive under high concurrency.
To address this, I implemented Group Commit. Instead of committing every update independently, the gateway groups multiple updates from the same tenant into a batch. It also coalesces multiple modifications to the same database row when possible. This significantly reduces the overhead associated with database round trips and repeated updates.
With 64 concurrent connections, PostgreSQL-backed throughput improved from approximately: 70 req/s → 877 req/s (12.5× improvement) .
Formal Verification with Lean 4
Making infrastructure faster is useful, but an AI gateway also needs to be correct. This becomes especially important when the gateway handles authentication, authorization, and financial accounting.
For example, consider the following requirements:
- A request must not be admitted if it violates the budget policy.
- An API key without administrative permissions must not be allowed to modify restricted limits.
- An optimized batch operation must preserve the semantics of the corresponding sequential operations.
These are not merely performance considerations. They are correctness properties. Traditional unit tests can check specific examples, but they cannot generally establish that a property holds for all possible inputs and states. This is where formal verification comes in.
Rust Kernels and Lean 4
One of the motivations behind h5i-app is to make formal verification more practical for real web applications. The framework encourages developers to implement important business logic as a separate Rust kernel. That kernel can then be translated into Lean 4, where mathematical properties of its behavior can be specified and proven.
For the AI gateway, this includes reasoning about admission decisions, permissions, and budget accounting. I also verified properties related to the batching optimization, including equivalence between optimized batched operations and their sequential counterparts under the modeled assumptions.
Try It Yourself
I have open-sourced both the AI gateway and its benchmark.
- AI Gateway: h5i-dev/h5i-ai-gateway
- Benchmark Scripts and Results: h5i-ai-gateway/bench
- Rust Web Framework: h5i-dev/h5i
