Introduction to TensorRT-LLM Optimization
TensorRT-LLM is NVIDIA's open-source library for optimizing Large Language Model inference on NVIDIA GPUs. Built on top of TensorRT, it provides a Python API to compile, optimize, and execute LLMs with state-of-the-art performance. Achieving sub-second inference — where the time-to-first-token (TTFT) and inter-token latency fall below one second — is critical for real-time applications such as chatbots, coding assistants, and interactive agents.
This tutorial walks through the practical techniques required to push TensorRT-LLM into sub-second territory. We will cover engine building, quantization, in-flight batching, KV cache management, paged attention, and deployment best practices. By the end, you will have a reproducible workflow you can adapt to your own models.
Why Sub-Second Inference Matters
Human perception of responsiveness is anchored around the one-second mark. When a model takes longer than a second to produce its first token, users perceive the system as sluggish. For agentic workflows where multiple LLM calls are chained together, latency compounds rapidly — a two-second per-call latency becomes twenty seconds across a ten-step pipeline.
Sub-second inference unlocks several capabilities:
- Real-time conversational UX: First-token latency under 200ms feels instantaneous.
- Higher throughput: Lower per-request latency means more concurrent users on the same hardware.
- Agentic loops: Multi-step reasoning, tool calling, and reflection become practical.
- Cost efficiency: Less GPU-seconds per request translates directly to lower cloud spend.
Prerequisites and Environment Setup
Before optimizing, ensure your environment is correctly configured. TensorRT-LLM requires an NVIDIA GPU with CUDA 12.x, the TensorRT-LLM Python package, and a model checkpoint. The examples below assume an NVIDIA A100 or H100, but the workflow applies to any modern NVIDIA datacenter GPU.
# Create a clean environment
conda create -n trt-llm python=3.10 -y
conda activate trt-llm
# Install TensorRT-LLM (matching your CUDA version)
pip install tensorrt-llm --extra-index-url https://pypi.nvidia.com
# Install supporting libraries
pip install transformers torch sentencepiece
# Verify installation
python -c "import tensorrt_llm; print(tensorrt_llm.__version__)"
For best results, use the official TensorRT-LLM container, which bundles compatible versions of CUDA, cuDNN, and TensorRT:
docker run --gpus all --rm -it \
-v $(pwd):/workspace \
nvcr.io/nvidia/tensorrt-llm:latest \
bash
Building an Optimized Engine
The core of TensorRT-LLM optimization is the engine build step. During this step, TensorRT fuses kernels, selects optimal algorithms, and bakes in configuration parameters. Choices made at build time have a larger impact on latency than runtime tuning, so it pays to configure carefully.
Converting a Hugging Face Model
TensorRT-LLM ships with conversion scripts for popular architectures (Llama, Mistral, GPT, Falcon, and others). The first step is to convert the Hugging Face checkpoint into the TensorRT-LLM format.
git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM/examples/llama
# Convert the model to the TRT-LLM checkpoint format
python convert_checkpoint.py \
--model_dir /models/llama-2-7b-hf \
--output_dir /models/llama-2-7b-trt \
--dtype float16
Building the Engine with Builder Flags
The build.py script exposes dozens of flags. The most impactful for latency are those controlling precision, attention implementation, and shape ranges.
python build.py \
--model_dir /models/llama-2-7b-trt \
--output_dir /engines/llama-2-7b \
--dtype float16 \
--use_gpt_attention_plugin float16 \
--use_gemm_plugin float16 \
--use_rmsnorm_plugin float16 \
--use_paged_context_kv_cache enable \
--max_batch_size 32 \
--max_input_len 1024 \
--max_output_len 256 \
--max_num_tokens 8192 \
--gather_generation_logits
Key flags explained:
--use_gpt_attention_plugin: Enables fused attention kernels, which are significantly faster than the native implementation.--use_paged_context_kv_cache: Enables paged KV cache, reducing memory fragmentation and allowing higher batch sizes.--max_num_tokens: Sets the maximum number of tokens the engine can process in a single batch. This controls the in-flight batching capacity.--gather_generation_logits: Required if you need the full logits output (e.g., for custom sampling). Disable for pure greedy decoding to save time.
Quantization for Latency Reduction
Quantization is the single most effective technique for reducing both latency and memory footprint. TensorRT-LLM supports INT8 SmoothQuant, INT4 AWQ, and FP8 on Hopper GPUs. FP8 is particularly attractive because it requires no calibration data and delivers near-FP16 quality.
FP8 Quantization on Hopper
FP8 requires an H100 or newer GPU. The conversion is nearly transparent — you simply change the dtype flag during checkpoint conversion and engine build.
# Convert with FP8
python convert_checkpoint.py \
--model_dir /models/llama-2-7b-hf \
--output_dir /models/llama-2-7b-fp8 \
--dtype float8
# Build the FP8 engine
python build.py \
--model_dir /models/llama-2-7b-fp8 \
--output_dir /engines/llama-2-7b-fp8 \
--dtype float16 \
--use_gpt_attention_plugin float16 \
--use_gemm_plugin float16 \
--use_rmsnorm_plugin float16 \
--max_batch_size 64 \
--max_input_len 1024 \
--max_output_len 256 \
--max_num_tokens 16384
FP8 typically reduces per-token latency by 30-50% compared to FP16 while doubling the achievable batch size, which is the key to sub-second performance under concurrent load.
INT4 AWQ for Maximum Compression
For memory-constrained deployments, INT4 AWQ offers the highest compression ratio. It requires a calibration step but produces engines that fit in a fraction of the VRAM.
# Quantize with AWQ
python ../quantization/quantize.py \
--model_dir /models/llama-2-7b-hf \
--output_dir /models/llama-2-7b-awq \
--dtype float16 \
--qformat int4_awq \
--calib_size 512 \
--calib_batch_size 32
# Build the AWQ engine
python build.py \
--model_dir /models/llama-2-7b-awq \
--output_dir /engines/llama-2-7b-awq \
--dtype float16 \
--use_gpt_attention_plugin float16 \
--use_gemm_plugin float16 \
--max_batch_size 128 \
--max_input_len 1024 \
--max_output_len 256 \
--max_num_tokens 16384
In-Flight Batching for Concurrent Requests
Static batching waits to fill a batch before processing, which inflates latency under low load. In-flight batching (also called continuous batching) inserts new requests into active batches at token boundaries, keeping the GPU saturated while minimizing per-request wait time.
The Triton Inference Server provides a first-class TensorRT-LLM backend with in-flight batching enabled by default. The following configuration enables it:
# config.pbtxt for the TensorRT-LLM backend
name: "tensorrt_llm"
backend: "tensorrtllm"
max_batch_size: 0
parameters [
{
key: "tokenizer_dir"
value: { string_value: "/models/llama-2-7b-hf" }
},
{
key: "engine_dir"
value: { string_value: "/engines/llama-2-7b-fp8" }
},
{
key: "max_tokens_in_paged_kv_cache"
value: { string_value: "8192" }
},
{
key: "max_attention_window_size"
value: { string_value: "1024" }
},
{
key: "kv_cache_free_gpu_mem_fraction"
value: { string_value: "0.85" }
},
{
key: "enable_trt_overlap"
value: { string_value: "true" }
}
]
dynamic_batching {
preferred_batch_size: [4, 8, 16, 32]
max_queue_delay_microseconds: 10000
preserve_ordering: true
}
The max_queue_delay_microseconds parameter is critical for sub-second latency. Setting it to 10ms ensures that requests are not held in the queue longer than necessary, while still allowing the scheduler to form reasonably sized batches.
Using the Python Runtime Directly
For cases where you need fine-grained control or want to embed TensorRT-LLM in a custom serving stack, you can use the Python runtime directly. The following example demonstrates a minimal generation loop with in-flight batching.
import torch
from tensorrt_llm.runtime import ModelRunner, ModelRunnerCpp
# Load the engine
runner = ModelRunnerCpp.from_dir(
engine_dir="/engines/llama-2-7b-fp8",
rank=0,
max_output_len=256,
)
# Prepare inputs
input_text = "Explain quantum computing in one sentence."
input_ids = tokenizer.encode(input_text, return_tensors="pt").cuda()
input_lengths = torch.tensor([input_ids.size(1)], dtype=torch.int32).cuda()
# Run generation
output_ids = runner.generate(
input_ids,
input_lengths,
max_new_tokens=64,
temperature=0.7,
top_p=0.9,
end_id=tokenizer.eos_token_id,
pad_id=tokenizer.pad_token_id,
)
# Decode
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)
For production serving, prefer the Triton backend over the raw Python runtime. Triton handles connection management, dynamic batching, and metrics collection, all of which are essential for maintaining sub-second latency under load.
KV Cache Tuning
The KV cache is the largest consumer of GPU memory during inference and directly affects how many concurrent requests the system can serve. Tuning the cache is a balancing act between memory usage, batch size, and context length.
The most important parameter is kv_cache_free_gpu_mem_fraction, which controls how much free GPU memory is allocated to the KV cache after loading the model weights. A value of 0.85 reserves 85% of remaining memory for the cache, leaving 15% for intermediate activations and overhead.
# Calculate optimal KV cache size
import torch
total_mem = torch.cuda.get_device_properties(0).total_memory
model_weights_mem = 14 * 1024**3 # ~14GB for a 7B FP8 model
free_mem = total_mem - model_weights_mem
kv_cache_mem = int(free_mem * 0.85)
# Estimate max concurrent sequences
# Each token in KV cache uses ~2 * num_layers * num_kv_heads * head_dim * 2 bytes
num_layers = 32
num_kv_heads = 32
head_dim = 128
bytes_per_token = 2 * num_layers * num_kv_heads * head_dim * 2 # FP16 KV
max_kv_tokens = kv_cache_mem // bytes_per_token
print(f"KV cache: {kv_cache_mem / 1024**3:.1f} GB")
print(f"Max KV tokens: {max_kv_tokens}")
print(f"Concurrent 1024-token contexts: {max_kv_tokens // 1024}")
Best Practices for Sub-Second Performance
Right-Size Your Context Window
Longer context windows dramatically increase both memory usage and attention computation. If your application only needs 1024 tokens of context, do not build the engine with a 4096 max input length. The attention kernel scales quadratically with sequence length, so halving the context can reduce TTFT by up to 4x.
Use Speculative Decoding
Speculative decoding uses a small draft model to propose tokens that the large model verifies in parallel. This can increase throughput by 2-3x for tasks with predictable output structure. TensorRT-LLM supports speculative decoding natively:
python build.py \
--model_dir /models/llama-2-7b-trt \
--output_dir /engines/llama-2-7b-speculative \
--dtype float16 \
--use_gpt_attention_plugin float16 \
--speculative_model /models/llama-2-160m-trt \
--speculative_model_type llama \
--max_draft_len 5
Enable Graph Rewriting and Kernel Fusion
Always enable the plugin flags during build. The fused RMSNorm, GEMM, and attention plugins eliminate kernel launch overhead and reduce memory bandwidth pressure. On small batch sizes, kernel launch overhead can dominate total latency, so fusion is essential for sub-second performance.
Profile and Iterate
Use the built-in profiler to identify bottlenecks. The Triton TensorRT-LLM backend exposes per-iteration metrics through Prometheus, including TTFT, inter-token latency, and queue time.
# Enable metrics in Triton
# In config.pbtxt, add:
parameters [
{
key: "enable_kv_cache_reuse"
value: { string_value: "true" }
}
]
# Query metrics
curl http://localhost:8002/metrics | grep trt_llm
Batch Size and Concurrency
Sub-second latency is achievable at low batch sizes with almost any configuration, but the real challenge is maintaining it under load. Profile your engine at the batch sizes you expect in production. If TTFT exceeds one second at batch size 16, consider FP8 quantization, reducing context length, or upgrading to a higher-memory GPU.
Conclusion
Achieving sub-second inference with TensorRT-LLM is a systematic process: build with fused plugins, quantize to FP8 or INT4, enable in-flight batching via Triton, tune the KV cache to your hardware, and right-size the context window for your application. Each optimization compounds — FP8 halves memory and latency, in-flight batching keeps the GPU saturated, and paged attention prevents fragmentation under concurrent load. By following the workflow in this tutorial and profiling iteratively against your real workload, you can deliver LLM inference that feels instantaneous to users while maximizing the return on your GPU investment.