From Model Weights to API Endpoint with TensorRT-LLM

Philip Kiely, Baseten, Pankaj Gupta, Baseten1:40:01 · Sept 2024 · 5,858 views
Thumbnail for From Model Weights to API Endpoint with TensorRT-LLM Watch on YouTube
TL;DR
  1. 1

    TensorRT-LLM adds optimized CUDA kernels, plugins, and in-flight batching to TensorRT for faster large language model inference on NVIDIA GPUs.

  2. 2

    An engine must be built around its target GPU, batch size, and input and output sequence lengths because those settings determine which kernels and optimizations are used.

  3. 3

    Benchmarking different concurrency and batch-size settings is necessary because throughput and latency do not change in a simple linear way.

Summary

Philip Kiely and Pankaj Gupta explain how TensorRT-LLM turns a model checkpoint into an engine that can run behind an API endpoint. They cover TensorRT's graph optimization model, TensorRT-LLM plugins for transformer operations, in-flight batching, quantization, tensor parallelism, and GPU-specific engine builds. The workshop uses TinyLlama 1.1B on cloud RTX 4090 and A10 GPUs so participants can complete the development loop quickly. The speakers convert a Hugging Face checkpoint, build an FP16-style engine, build an FP8 engine with calibration, run inference, and compare engine sizes. They then show how Baseten's Truss packages the same process in YAML and deploys it through an API. The final section focuses on benchmarking, including time to first token, per-request tokens per second, total throughput, concurrency, and percentile latency. Production deployment adds problems around large engine files, cold starts, GPU matching, batching, and autoscaling.

Key ideas
03:32

TensorRT optimizes computation graphs before executing them on NVIDIA GPUs

Pankaj Gupta describes machine learning models as computation graphs through which data flows. TensorRT takes a graph representation, finds patterns, optimizes them, and creates an executable engine. It supports C, C++, and Python, although the workshop uses Python. TensorRT also has a plugin mechanism for operations it cannot infer or optimize automatically, such as flash attention. TensorRT-LLM uses that mechanism to provide specialized plugins and CUDA kernels for large language model workloads. The result is an engine that can execute the graph with lower latency and better GPU efficiency than a direct graph execution path.

06:05

TensorRT-LLM adds language-model features that TensorRT does not provide by itself

TensorRT is primarily a graph executor, while TensorRT-LLM adds features specific to transformer inference. The speakers point to optimized attention and flash-attention plugins, LoRA support, Medusa heads, speculative decoding, paged attention, and in-flight batching. Pankaj calls in-flight batching essential for serving language models because requests arrive and generate tokens at different times. TensorRT-LLM also lets builders define optimization profiles around input length, output length, and batch size. Those values matter because the system can select different kernels for different matrix and sequence sizes. The trade-off is that engine compilation can take hours when many sizes are included.

10:37

A TensorRT-LLM engine is tailored to its deployment environment

The speakers compare engine optimization to having a suit made to a person's measurements. Builders need to know the production sequence lengths, batch sizes, and GPU type in advance. An engine built for an A10 is intended to run on an A10, while an engine built for an H100 should be built again for that GPU or its MIG configuration. TensorRT-LLM also creates infrastructure concerns because engines can be large and can make autoscaling cold starts slow. There are modes that allow an engine built for one GPU to run on another, but Pankaj says those modes sacrifice optimization and are not what Baseten normally uses.

18:32

GPU architecture and precision affect the gains from TensorRT-LLM

TensorRT compiles models to CUDA instructions, so newer GPU architecture features can produce gains beyond the raw increase in FLOPS or memory bandwidth. The speakers describe H100s as especially useful for smaller models on MIG partitions because those models may not need the full GPU memory. They also discuss FP8, which H100 and L4 GPUs support while A100 and A10 GPUs do not. Pankaj says FP8 can be applied to weights and to the KV cache. Weight quantization usually causes little quality loss, while KV-cache quantization has historically been more difficult. In their experience, FP8 makes KV-cache quantization practical, although calibration data still affects the selected ranges.

29:29

The manual build starts with checkpoint conversion, engine compilation, and a direct inference test

The workshop uses TinyLlama 1.1B to keep the development loop short, rather than waiting for a large production model to compile. The process runs on a cloud GPU, installs TensorRT-LLM, downloads the Hugging Face model, and converts its checkpoint into the format expected by TensorRT-LLM. That conversion can also handle some quantization paths and divide weights for tensor parallelism. The build command then creates and serializes a TensorRT engine. The speakers describe the engine as similar to a binary or shared library that the TensorRT-LLM runtime can load. They run the finished engine with a standard example prompt and receive generated text.

56:17

FP8 can reduce engine size while preserving similar output quality

For the second engine build, the speakers use the NVIDIA AMMO library and specify FP8 quantization for the model weights and KV cache. They explain that FP8's exponent gives it more useful dynamic range than a linear INT8 representation, which helps with smaller values. Calibration data is used to estimate minimum and maximum ranges for the transformer layers. The FP8 engine is about 1.2 GB, compared with about 2 GB for the earlier engine in their TinyLlama example. They run the same prompt and observe output that is very similar to the previous result. Pankaj says FP8 quality is difficult to distinguish from the higher-precision output in their tests.

01:03:44

Truss packages TensorRT-LLM configuration and deployment behind an API

Baseten's open-source Truss framework provides a packaging and serving layer around model code. The speakers describe it as a configuration file plus Python code, with support for TensorRT, TensorRT-LLM, vLLM, Triton, Transformers, Diffusers, and other model types. In this workshop, they do not write a model-server implementation. Instead, they configure the model in YAML with the model source, GPU, sequence lengths, batch size, and quantization settings. Their example uses an input and output sequence length of 2,000 tokens, a batch size of 64, and INT8 quantization on an A10. Truss then builds and deploys the model so it can be called through an HTTP API.

01:17:09

Benchmarking must measure the actual workload because batch-size effects are irregular

The speakers recommend measuring time to first token, tokens per second per request, total tokens per second, and latency percentiles such as the 50th, 90th, 95th, and 99th. They show that increasing batch size does not produce a smooth performance curve. In one Mistral example, doubling the batch size from 32 to 64 barely changes time to first token, while increasing it from 64 to 128 causes a large jump. Higher batch sizes can improve total throughput, but they can also slow individual requests and consume more KV-cache memory. In the workshop benchmark, concurrency of 32 produces about 5,000 total tokens per second for TinyLlama, while a later run at concurrency 64 reaches about 7,000. These results depend on the deployed engine and its maximum batch size.

01:30:00

Production autoscaling has to account for engine size, GPU matching, and long-lived streaming requests

The production serving path uses Triton Inference Server or a related server layer to load the TensorRT engine and handle requests. The speakers describe challenges with horizontal scaling, including routing, batching, GPU availability, large image and engine sizes, and matching new nodes to the engine's GPU requirements. Baseten uses a modified version of KNative because ordinary KNative assumptions target short microservice requests, while model requests can stream for a long time. Pods scale out as active capacity fills and scale down toward zero when traffic drops. Requests can remain queued during scale-up. The speakers say fast cold starts are required for scale-to-zero to work, with stated goals of under 10 seconds for 10 GB models and under a minute for 50 GB models.

"When you quantize both the KV cache and the weights to FP8, you get that extra unlock that your computation is also faster."57:48
Who should watch
  • You are choosing between TensorRT-LLM, vLLM, or another serving stack and need to understand the performance and engineering trade-offs.
  • You want to build a GPU-specific TensorRT-LLM engine, test FP8 or INT8 quantization, and measure latency against throughput before production.
  • Your model already runs, but large engines, cold starts, batching, and GPU-specific deployment make autoscaling difficult.