DEV Community

Benito August
Benito August

Posted on

Under 3ms Latency Building a High-Performance LLM Gateway in Rust

Image generated with Gemini Nano Banana

Scaling Large Language Model (LLM) integrations in production introduces three critical bottlenecks: high cloud API token costs, latency overheads, and data privacy leaks. When you use Retrieval-Augmented Generation (RAG) or feed long chat histories back to models like Gemini, Claude, or GPT-4, you are not only paying for redundant context, but you are also waiting for the network round-trip.

Traditionally, developers turn to Python-based microservices for LLM orchestration and preprocessing. However, Python’s interpreter overhead, lack of true concurrency, and high memory footprint make it a bottleneck when running local deep-learning models (like BERT-based NER or MiniLM sentence embedders) at high throughput.

To solve this, we designed NLProxy, a local-first, high-performance LLM gateway written in Rust. In this post, we’ll dive into how we leveraged the Rust systems programming ecosystem (ort, candle, tokio, regex, and linfa-clustering) to build a gateway capable of handling 5000+ requests per second with sub-3ms latency.


The Core Concept: Local-First Hybrid LLM Gateway

A local hybrid gateway sits between your internal application threads and public cloud APIs. Before any prompt egresses your network, the gateway intercepts it to perform:

  1. PII Masking (Shielding): Local Named Entity Recognition (NER) to redact names, emails, and IPs.
  2. Semantic Prompt Compression: Clustering sentences and discarding redundancies.
  3. Prompt Firewall Check: Detecting jailbreaks, SQL injections, or ReDoS payloads in $O(n)$ linear time.
  4. Post-Generation Verification: Using local Natural Language Inference (NLI) to flag hallucinations before returning the answer.

Doing all of this locally under 3 milliseconds is impossible in Python. Here is how Rust makes it achievable.


1. Zero-Copy PII Shielding with candle

Redacting Personally Identifiable Information (PII) before it reaches third-party servers is a strict compliance requirement under GDPR and HIPAA. In NLProxy, we use a hybrid approach: deterministic regular expressions combined with a local Named Entity Recognition (NER) model running on CPU.

For the NER engine, we leverage Hugging Face's candle, a minimalist machine learning framework for Rust. Unlike Python’s heavy PyTorch bindings, candle lets us run lightweight BERT models natively without virtual environments or massive dynamically linked libraries.

Here is a look at how NLProxy's NerEngine loads a safe tensor BERT model and predicts entity spans:

use candle_core::{Device, Tensor};
use candle_transformers::models::bert::{BertModel, Config};
use tokenizers::Tokenizer;
use candle_nn::VarBuilder;

pub struct NerEngine {
    model: BertModel,
    tokenizer: Tokenizer,
    classifier_weight: Tensor,
    classifier_bias: Tensor,
    device: Device,
}

impl NerEngine {
    pub fn load(models_dir: &std::path::Path) -> anyhow::Result<Self> {
        let device = Device::Cpu;
        let folder = models_dir.join("bert-base-NER");
        let config_str = std::fs::read_to_string(folder.join("config.json"))?;
        let config: Config = serde_json::from_str(&config_str)?;
        let tokenizer = Tokenizer::from_file(folder.join("tokenizer.json"))?;

        let vb = unsafe {
            VarBuilder::from_mmaped_safetensors(&[folder.join("model.safetensors")], candle_core::DType::F32, &device)?
        };

        let model = BertModel::load(vb.pp("bert"), &config)?;
        let classifier_weight = vb.get((9, 768), "classifier.weight")?;
        let classifier_bias = vb.get(9, "classifier.bias")?;

        Ok(NerEngine { model, tokenizer, classifier_weight, classifier_bias, device })
    }
}
Enter fullscreen mode Exit fullscreen mode

By memory-mapping safetensors (VarBuilder::from_mmaped_safetensors), we eliminate the startup overhead and minimize RAM usage. The system redacts PII into random token hashes (e.g. __PROT_82739182__) stored locally in memory, restoring them instantly when the LLM returns the output.


2. Semantic Prompt Compression with linfa and kodama

To reduce input token fees, we compress prompts by clustering similar sentences and keeping only the most representative ones. We run a sentence-embedding model local ONNX via the ort crate (ONNX Runtime bindings) to obtain dense vector arrays (Array2<f32> from the ndarray crate).

Once embeddings are computed, we can choose between:

  • KMeans Clustering (using the linfa-clustering crate) for massive contexts.
  • Hierarchical Ward Linkage (using the kodama crate) for precise structural grouping.

Here is a snippet showing how NLProxy's semantic compressor handles KMeans sentence reduction locally in Rust:

use linfa_clustering::KMeans;
use linfa::traits::{Fit, Predict};
use ndarray::Array2;

fn cluster_kmeans(
    sentences: Vec<String>,
    embeddings: &Array2<f32>,
    original_indices: Vec<usize>,
    aggressiveness: f32,
) -> (Vec<String>, Vec<usize>) {
    let n = sentences.len();
    let n_clusters = (n as f32 * (1.0 - aggressiveness)).max(1.0).min(n as f32) as usize;

    if n_clusters >= n {
        return (sentences, original_indices);
    }

    let kmeans = KMeans::params(n_clusters)
        .max_n_iterations(100)
        .tolerance(1e-4);

    let dataset = linfa::Dataset::from(embeddings.clone());
    let model = kmeans.fit(&dataset).expect("KMeans clustering failed");
    let labels = model.predict(&dataset);

    let mut selected_sentences = Vec::new();
    let mut selected_indices = Vec::new();

    for cluster_idx in 0..n_clusters {
        let mut cluster_member_indices = Vec::new();
        for (i, &label) in labels.iter().enumerate() {
            if label == cluster_idx {
                cluster_member_indices.push(i);
            }
        }

        if cluster_member_indices.is_empty() { continue; }

        let centroid = model.centroids().index_axis(ndarray::Axis(0), cluster_idx);
        let mut best_idx = cluster_member_indices[0];
        let mut min_dist = f32::MAX;

        for &idx in &cluster_member_indices {
            let emb = embeddings.index_axis(ndarray::Axis(0), idx);
            let dist = emb.iter().zip(centroid.iter())
                .map(|(a, b)| (a - b).powi(2))
                .sum::<f32>()
                .sqrt();

            if dist < min_dist {
                min_dist = dist;
                best_idx = idx;
            }
        }

        selected_sentences.push(sentences[best_idx].clone());
        selected_indices.push(original_indices[best_idx]);
    }

    // Sort selected sentences back to their original document order
    let mut combined: Vec<_> = selected_sentences.into_iter().zip(selected_indices.into_iter()).collect();
    combined.sort_by_key(|&(_, idx)| idx);
    combined.into_iter().unzip()
}
Enter fullscreen mode Exit fullscreen mode

Executing this in Rust takes a fraction of a millisecond. We get a 40% to 55% reduction in token count without losing the system instructions or core semantic context.


3. ReDoS-Immune Firewall Engine

LLM application firewalls must defend against prompt injection attacks (like jailbreaks) and system denial-of-service attempts. Python and JavaScript regular expression engines are prone to ReDoS (Regular Expression Denial of Service) due to backtracking. An attacker can construct a complex query that causes the CPU to lock up, taking down your gateway.

Rust’s regex crate solves this fundamentally. It guarantees linear-time search $O(m \times n)$ with respect to the input length by using Finite Automata (DFA/NFA) representation, avoiding backtracking entirely.

NLProxy wraps these rules into deterministic DFA checks, and overlays semantic classification using cosine similarity against an attack vector database. The CPU overhead is negligible (measured in micro-seconds), keeping the firewall execution invisible to user request threads.


4. Hallucination Guardrails with Local NLI

Before sending the LLM output back to the user, we evaluate it against the original prompt or source document to check for logical contradictions. Running another cloud LLM call to evaluate the answer is slow and expensive.

Instead, NLProxy embeds a lightweight cross-lingual XLM-RoBERTa model set up for Natural Language Inference (NLI) classification (Entailment vs Contradiction). Running this locally via candle classification layer allows us to identify contradictions on the fly:

use candle_transformers::models::xlm_roberta::XLMRobertaForSequenceClassification;

pub fn predict_contradiction(&self, premise: &str, hypothesis: &str) -> anyhow::Result<f32> {
    let tokens = self.tokenizer.encode((premise, hypothesis), true).map_err(anyhow::Error::msg)?;
    let token_ids = tokens.get_ids();
    let input_ids = Tensor::new(token_ids, &self.device)?.unsqueeze(0)?;
    let attention_mask = Tensor::new(vec![1u32; token_ids.len()], &self.device)?.unsqueeze(0)?;
    let token_type_ids = Tensor::new(vec![0u32; token_ids.len()], &self.device)?.unsqueeze(0)?;

    let logits = self.model.forward(&input_ids, &attention_mask, &token_type_ids)?;
    let logits_vec = logits.squeeze(0)?.to_vec1::<f32>()?;

    let max_val = logits_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let sum_exp: f32 = logits_vec.iter().map(|&val| (val - max_val).exp()).sum();
    let contradiction_prob = (logits_vec[0] - max_val).exp() / sum_exp; // index 0 matches Contradiction

    Ok(contradiction_prob)
}
Enter fullscreen mode Exit fullscreen mode

If the contradiction probability exceeds a threshold (e.g. 70%), the gateway intercepts the response, blocks the output, or corrects placeholders natively before the user ever sees it.


Real-World Performance & Benchmarks

Running the entire unified pipeline (PII Shielding, Semantic Compression, Firewall, LLM Call, and Post-Verification) locally in Rust vs. Python yields massive efficiency gains:

Metric Python (FastAPI + PyTorch) Rust Core (NLProxy) Improvement
Max Throughput ~180 req/sec 5000+ req/sec 27.7x
Pre-LLM Latency ~65 ms < 3 ms 21.6x
RAM Footprint ~4.2 GB ~1.2 GB 3.5x
ReDoS Immunity Vulnerable (PCRE) 100% Immune Architectural

Conclusion

By writing the core gateway logic in Rust, NLProxy shows that systems programming is no longer just for operating systems or databases. For AI engineering, building compiled, native middleware on top of crates like candle and ort lets you run local models, guarantee absolute privacy, and slash cloud token bills without degrading your API response times.

If you are looking to build secure, low-latency, and cost-effective AI features, it’s time to move your integration middleware from Python to Rust.


NLProxy is an open-core project. Check out the repository on GitHub and start building high-performance AI middleware.

Top comments (0)