How to Use ONNX Runtime for LLM Inference

Introduction to ONNX Runtime for LLM Inference

Large Language Models (LLMs) have transformed how developers build intelligent applications, but running them efficiently in production remains a significant challenge. ONNX Runtime, originally designed for traditional machine learning models, has evolved into a powerful inference engine capable of handling transformer-based architectures. By leveraging ONNX Runtime for LLM inference, developers can achieve cross-platform compatibility, hardware acceleration, and consistent performance across diverse deployment environments.

This tutorial walks you through everything you need to know to get started with ONNX Runtime for LLM inference, from understanding the fundamentals to deploying optimized models in production.

What Is ONNX Runtime?

ONNX Runtime is an open-source, cross-platform inference accelerator built by Microsoft. It implements the Open Neural Network Exchange (ONNX) format, an open standard for representing machine learning models. The runtime provides a unified API that can execute models across CPUs, GPUs, and specialized accelerators without requiring framework-specific dependencies at inference time.

For LLMs specifically, ONNX Runtime includes dedicated support through the onnxruntime-genai package, which handles the unique requirements of generative text models such as tokenization, KV-cache management, sampling strategies, and iterative decoding loops.

Key Components

• ONNX Runtime Core: The base execution engine that runs ONNX graph computations.
• Execution Providers: Pluggable backends like CUDA, TensorRT, DirectML, OpenVINO, and CoreML that enable hardware acceleration.
• ONNX Runtime GenAI: A specialized library built on top of the core runtime that provides LLM-specific APIs for generation, chat, and streaming.
• Model Optimizer: Tools for converting and optimizing models from PyTorch, Hugging Face, and other frameworks into ONNX format.

Why Use ONNX Runtime for LLM Inference?

Choosing the right inference engine is critical for production LLM deployments. ONNX Runtime offers several compelling advantages that make it an excellent choice for many scenarios.

Cross-Platform Portability

Once a model is converted to ONNX format, it can run identically on Windows, Linux, macOS, and even mobile platforms. This eliminates the need to maintain separate deployment pipelines for different operating systems or hardware configurations.

Hardware Acceleration

ONNX Runtime supports a wide range of execution providers, allowing you to tap into hardware acceleration with minimal code changes. The same model can run on an NVIDIA GPU using CUDA or TensorRT, an AMD GPU using ROCm, an Intel CPU using OpenVINO, or an Apple Silicon Mac using CoreML.

Optimized Performance

The runtime includes graph optimizations such as operator fusion, constant folding, and memory layout transformations. For LLMs, it supports techniques like KV-cache management, grouped-query attention, and quantization-aware execution that significantly reduce latency and memory consumption.

Production Readiness

ONNX Runtime is battle-tested in production at Microsoft and other large organizations. It provides robust error handling, thread safety, telemetry hooks, and consistent behavior across versions, making it suitable for enterprise deployments.

Setting Up Your Environment

Before diving into code, you need to install the necessary packages. The setup differs depending on whether you want CPU-only inference or GPU acceleration.

Installing for CPU Inference

pip install onnxruntime-genai
pip install numpy
pip install transformers

Installing for GPU Inference

pip install onnxruntime-genai-cuda
pip install numpy
pip install transformers

For GPU setups, ensure you have the appropriate CUDA toolkit and cuDNN libraries installed on your system. The onnxruntime-genai-cuda package bundles compatible runtime libraries, but matching your system CUDA version is important for stability.

Converting Models to ONNX Format

The first step in using ONNX Runtime for LLM inference is obtaining an ONNX-formatted model. There are two primary approaches: using pre-converted models or converting models yourself.

Using Pre-Converted Models

Many popular LLMs are already available in ONNX format on Hugging Face. Models from the onnx-community organization are specifically optimized for ONNX Runtime. You can browse and download these directly.

from huggingface_hub import snapshot_download

model_path = snapshot_download(
    repo_id="onnx-community/Llama-3.2-1B-Instruct-onnx",
    local_dir="./llama-onnx"
)
print(f"Model downloaded to: {model_path}")

Converting from Hugging Face

If a pre-converted model is not available, you can convert a Hugging Face model yourself using the optimum library, which provides ONNX export utilities.

pip install optimum[onnxruntime]

from optimum.onnxruntime import ORTModelForCausalLM
from transformers import AutoTokenizer

model_id = "microsoft/phi-2"

# Export the model to ONNX format
model = ORTModelForCausalLM.from_pretrained(
    model_id,
    export=True,
    provider="CPUExecutionProvider"
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# Save the converted model
model.save_pretrained("./phi-2-onnx")
tokenizer.save_pretrained("./phi-2-onnx")
print("Model converted and saved successfully.")

For larger models, you may want to export with specific precision settings. The following example shows how to export with FP16 precision for GPU deployment.

from optimum.onnxruntime import ORTModelForCausalLM

model = ORTModelForCausalLM.from_pretrained(
    "microsoft/phi-2",
    export=True,
    provider="CUDAExecutionProvider",
    file_name="model_fp16.onnx",
    from_transformers=True
)
model.save_pretrained("./phi-2-onnx-fp16")

Running Basic LLM Inference

Once you have an ONNX-formatted model, you can use the onnxruntime-genai API to run inference. The following example demonstrates a complete text generation pipeline.

Simple Text Generation

import onnxruntime_genai as og

# Load the model
model_path = "./llama-onnx"
model = og.Model(model_path)

# Create a tokenizer
tokenizer = og.Tokenizer(model)

# Define the prompt
prompt = "Explain the concept of recursion in programming."

# Encode the prompt
input_tokens = tokenizer.encode(prompt)

# Configure generation parameters
search_options = {
    "max_length": 256,
    "temperature": 0.7,
    "top_p": 0.9,
    "do_sample": True
}

# Create a generator
generator = og.Generator(model, search_options)
generator.append_tokens(input_tokens)

# Generate tokens iteratively
print("Generated response:")
print(prompt, end="", flush=True)

while not generator.is_done():
    generator.generate_next_token()
    new_token = generator.get_next_tokens()
    new_text = tokenizer.decode(new_token)
    print(new_text, end="", flush=True)

print()

Chat-Based Generation

For instruction-tuned models, you should use a chat template to format conversations properly. This ensures the model receives input in the format it was trained on.

import onnxruntime_genai as og

model = og.Model("./llama-onnx")
tokenizer = og.Tokenizer(model)

# Build a conversation
messages = [
    {"role": "system", "content": "You are a helpful coding assistant."},
    {"role": "user", "content": "Write a Python function to reverse a linked list."}
]

# Apply the chat template
chat_input = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True
)

# Generate response
search_options = {
    "max_length": 512,
    "temperature": 0.3,
    "top_p": 0.9
}

generator = og.Generator(model, search_options)
generator.append_tokens(chat_input)

response_tokens = []
while not generator.is_done():
    generator.generate_next_token()
    response_tokens.append(generator.get_next_tokens()[0])

response = tokenizer.decode(response_tokens)
print("Assistant:", response)

Streaming Responses

In production applications, streaming tokens to the user as they are generated dramatically improves perceived performance. ONNX Runtime GenAI supports this naturally through its iterative generation loop.

import onnxruntime_genai as og

class StreamingGenerator:
    def __init__(self, model_path):
        self.model = og.Model(model_path)
        self.tokenizer = og.Tokenizer(self.model)

    def generate_stream(self, prompt, max_length=256, temperature=0.7):
        input_tokens = self.tokenizer.encode(prompt)

        search_options = {
            "max_length": max_length,
            "temperature": temperature
        }

        generator = og.Generator(self.model, search_options)
        generator.append_tokens(input_tokens)

        while not generator.is_done():
            generator.generate_next_token()
            token = generator.get_next_tokens()
            chunk = self.tokenizer.decode(token)
            yield chunk

# Usage example
streamer = StreamingGenerator("./llama-onnx")
prompt = "Write a short story about a robot learning to paint."

print("Streaming output:")
for chunk in streamer.generate_stream(prompt):
    print(chunk, end="", flush=True)
print()

Using Execution Providers for Acceleration

One of the most powerful features of ONNX Runtime is its execution provider system. You can switch between hardware backends with minimal code changes.

GPU Acceleration with CUDA

import onnxruntime_genai as og

# The provider is typically configured in the model's config.json
# or can be specified when loading the model
model = og.Model("./llama-onnx", device="cuda")

tokenizer = og.Tokenizer(model)
prompt = "What are the benefits of renewable energy?"
input_tokens = tokenizer.encode(prompt)

search_options = {"max_length": 200, "temperature": 0.5}
generator = og.Generator(model, search_options)
generator.append_tokens(input_tokens)

while not generator.is_done():
    generator.generate_next_token()

output = tokenizer.decode(generator.get_sequence())
print(output)

DirectML for Windows GPU

On Windows systems with various GPU vendors (NVIDIA, AMD, Intel), DirectML provides a unified acceleration path.

import onnxruntime_genai as og

# Use DirectML provider (device index 0 is the default GPU)
model = og.Model("./llama-onnx", device="dml", device_id=0)

tokenizer = og.Tokenizer(model)
input_tokens = tokenizer.encode("Summarize the plot of Romeo and Juliet.")

search_options = {"max_length": 300}
generator = og.Generator(model, search_options)
generator.append_tokens(input_tokens)

while not generator.is_done():
    generator.generate_next_token()

print(tokenizer.decode(generator.get_sequence()))

Quantization for Reduced Memory and Faster Inference

Quantization reduces model precision from FP32 to INT8 or INT4, dramatically decreasing memory usage and improving inference speed with minimal accuracy loss. ONNX Runtime provides robust quantization tooling.

Dynamic Quantization

from onnxruntime.quantization import quantize_dynamic, QuantType

model_input = "./llama-onnx/model.onnx"
model_output = "./llama-onnx/model_int8.onnx"

quantize_dynamic(
    model_input,
    model_output,
    weight_type=QuantType.QInt8
)
print("Dynamic quantization complete.")

Using Pre-Quantized Models

Many models on Hugging Face are already available in quantized ONNX formats. Using INT4 quantized models can reduce memory requirements by up to 75% compared to FP16.

from huggingface_hub import snapshot_download

# Download an INT4 quantized model
model_path = snapshot_download(
    repo_id="onnx-community/Llama-3.2-1B-Instruct-onnx",
    local_dir="./llama-int4",
    allow_patterns=["*int4*"]
)

import onnxruntime_genai as og

model = og.Model(model_path)
tokenizer = og.Tokenizer(model)

input_tokens = tokenizer.encode("What is the capital of France?")
search_options = {"max_length": 100}
generator = og.Generator(model, search_options)
generator.append_tokens(input_tokens)

while not generator.is_done():
    generator.generate_next_token()

print(tokenizer.decode(generator.get_sequence()))

Batch Processing for Throughput

When serving multiple requests simultaneously, batch processing can significantly improve throughput. ONNX Runtime supports batched generation, though it requires careful management of padding and attention masks.

import onnxruntime_genai as og
import numpy as np

model = og.Model("./llama-onnx")
tokenizer = og.Tokenizer(model)

prompts = [
    "What is machine learning?",
    "Explain quantum computing briefly.",
    "How does photosynthesis work?"
]

# Tokenize all prompts
all_input_tokens = [tokenizer.encode(p) for p in prompts]

search_options = {"max_length": 128, "temperature": 0.5}

# Process each prompt (true batching support depends on model export)
results = []
for i, input_tokens in enumerate(all_input_tokens):
    generator = og.Generator(model, search_options)
    generator.append_tokens(input_tokens)

    while not generator.is_done():
        generator.generate_next_token()

    output = tokenizer.decode(generator.get_sequence())
    results.append(output)
    print(f"Result {i+1}: {output}\n")

Best Practices

Choose the Right Model Size

Select the smallest model that meets your quality requirements. A 1B parameter model often performs adequately for simple tasks while being dramatically faster and cheaper to run than a 7B or 13B model.

Optimize Your Prompt

Shorter, well-structured prompts reduce token processing time. Use clear system prompts and concise instructions to minimize unnecessary computation during both encoding and generation.

Manage KV-Cache Efficiently

The KV-cache stores previously computed attention keys and values, avoiding redundant computation during autoregressive generation. ONNX Runtime handles this internally, but be aware that longer conversations consume more memory. Consider implementing conversation length limits or summarization strategies for extended interactions.

Use Appropriate Precision

• FP32: Maximum accuracy, highest memory usage. Use only for development or when accuracy is critical.
• FP16: Good balance of accuracy and performance. Recommended for most GPU deployments.
• INT8: Significant speedup with minimal accuracy loss. Great for CPU deployment.
• INT4: Maximum compression. Best for edge devices or when memory is severely constrained.

Profile and Benchmark

Always measure performance in your specific deployment environment. Use the built-in profiling tools to identify bottlenecks.

import onnxruntime as ort

# Enable tracing for profiling
sess_options = ort.SessionOptions()
sess_options.enable_profiling = True
sess_options.profile_file_prefix = "onnx_profile"

# Use these options when creating your session
# Profile data will be saved as a JSON file
# View it with chrome://tracing or perfetto.dev

Handle Errors Gracefully

import onnxruntime_genai as og

def safe_generate(model_path, prompt, max_length=256):
    try:
        model = og.Model(model_path)
        tokenizer = og.Tokenizer(model)
        input_tokens = tokenizer.encode(prompt)

        if len(input_tokens) > max_length:
            raise ValueError("Input prompt exceeds maximum length")

        search_options = {"max_length": max_length}
        generator = og.Generator(model, search_options)
        generator.append_tokens(input_tokens)

        while not generator.is_done():
            generator.generate_next_token()

        return tokenizer.decode(generator.get_sequence())

    except Exception as e:
        print(f"Generation failed: {e}")
        return None

Building a Simple API Server

To serve your ONNX Runtime LLM in production, you can wrap it in a lightweight web API. Here is a minimal example using FastAPI with streaming support.

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import onnxruntime_genai as og
import json

app = FastAPI()

# Load model once at startup
model = og.Model("./llama-onnx")
tokenizer = og.Tokenizer(model)

class GenerateRequest(BaseModel):
    prompt: str
    max_length: int = 256
    temperature: float = 0.7

@app.post("/generate")
def generate(request: GenerateRequest):
    input_tokens = tokenizer.encode(request.prompt)
    search_options = {
        "max_length": request.max_length,
        "temperature": request.temperature
    }

    generator = og.Generator(model, search_options)
    generator.append_tokens(input_tokens)

    result = []
    while not generator.is_done():
        generator.generate_next_token()
        token = generator.get_next_tokens()
        chunk = tokenizer.decode(token)
        result.append(chunk)

    return {"response": "".join(result)}

@app.post("/stream")
def stream_generate(request: GenerateRequest):
    def event_stream():
        input_tokens = tokenizer.encode(request.prompt)
        search_options = {
            "max_length": request.max_length,
            "temperature": request.temperature
        }

        generator = og.Generator(model, search_options)
        generator.append_tokens(input_tokens)

        while not generator.is_done():
            generator.generate_next_token()
            token = generator.get_next_tokens()
            chunk = tokenizer.decode(token)
            yield f"data: {json.dumps({'token': chunk})}\n\n"

        yield f"data: {json.dumps({'done': True})}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")

# Run with: uvicorn server:app --host 0.0.0.0 --port 8000

Conclusion

ONNX Runtime provides a powerful, flexible, and production-ready solution for LLM inference across diverse hardware platforms. By converting your models to ONNX format, leveraging execution providers for hardware acceleration, applying quantization for memory efficiency, and following best practices for prompt engineering and error handling, you can build robust LLM applications that perform well in real-world deployments. The cross-platform nature of ONNX Runtime means you can develop on one platform and deploy on another with confidence, while the growing ecosystem of pre-converted models on Hugging Face makes getting started easier than ever. Whether you are building a chatbot, a code assistant, or a document summarization tool, ONNX Runtime offers the performance and reliability needed to ship LLM-powered features to production.