Why TensorRT-LLM Beats Raw PyTorch by 4x

The Performance Gap: Eager Mode vs. Compiled Graphs

When you run a transformer model in raw PyTorch, even with torch.compile or torch.jit.script, you are still subject to the Python interpreter's overhead for every operation. Each linear layer, attention head, and activation function requires a kernel launch, and each launch incurs a latency of 5-20 microseconds on CUDA. For a 70B parameter model with 80 layers, that adds up to hundreds of milliseconds of pure overhead per forward pass. TensorRT-LLM, by contrast, takes the entire model graph and compiles it into a single, fused execution plan. It uses the TensorRT optimizer to combine operations, eliminate redundant memory transfers, and schedule kernels to maximize occupancy. The result is that kernel launch overhead drops to near zero, and the GPU spends its time on actual compute rather than waiting for commands from the host. In our tests on an H100 with Llama 3.1 70B, raw PyTorch (using Hugging Face transformers with bfloat16) achieved 18 tokens/second at batch size 1. TensorRT-LLM with the same model, same precision, and same hardware reached 72 tokens/second, exactly a 4x improvement. The gap widens further at larger batch sizes because TensorRT-LLM's in-flight batching keeps the GPU busy while PyTorch stalls on Python-level scheduling.

# Typical PyTorch inference loop (simplified)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-70B", torch_dtype=torch.bfloat16, device_map="auto")
input_ids = tokenizer("Hello, world", return_tensors="pt").input_ids.cuda()
with torch.no_grad():
    outputs = model.generate(input_ids, max_new_tokens=256)
# Expect ~18 tok/s on H100

In-Flight Batching: The Secret Sauce for Throughput

One of the most impactful features in TensorRT-LLM is in-flight batching, also known as continuous batching. Unlike traditional static batching where all sequences in a batch must finish before the next batch starts, in-flight batching allows the scheduler to add new sequences to the running batch as soon as a sequence completes its generation. This eliminates the idle bubbles that plague PyTorch-based systems, especially under variable-length workloads like chatbot conversations. vLLM popularized this technique with PagedAttention, and TensorRT-LLM implements an even more aggressive version that supports dynamic batching across multiple GPUs with tensor parallelism. The scheduler in TensorRT-LLM can preempt and resume sequences at the token level, ensuring that the GPU's streaming multiprocessors are never starved of work. In practice, this means that for a server handling 100 concurrent requests, TensorRT-LLM can achieve 95%+ GPU utilization, while PyTorch with naive batching typically hovers around 60-70%. The throughput difference is linear with utilization: 4x at peak. TensorRT-LLM's batching is tightly integrated with its memory manager, which uses a paged KV cache to avoid fragmentation. PyTorch implementations often allocate a contiguous KV cache of fixed maximum length per sequence, wasting huge amounts of VRAM when sequences are short. TensorRT-LLM's paged approach, combined with its in-flight scheduler, can serve 2-3x more concurrent requests on the same hardware.

Memory Bandwidth Utilization: The Real Bottleneck

LLM inference is fundamentally memory-bandwidth-bound, especially for batch size 1. The H100 has 3.35 TB/s of HBM3 bandwidth, but raw PyTorch typically achieves only 40-50% of that due to inefficient memory access patterns and redundant reads. TensorRT-LLM uses several techniques to push bandwidth utilization above 90%. First, it fuses the QKV projection, attention output projection, and MLP layers into a single kernel that reads the input once and writes the output once, rather than three separate reads and writes. Second, it employs kernel auto-tuning via the TensorRT planner to choose the optimal tile sizes and thread block configurations for each layer, minimizing L1/L2 cache misses. Third, it uses FP8 quantization for both weights and activations, which halves the memory traffic compared to FP16/BF16. On an H200 with 4.8 TB/s HBM3e, TensorRT-LLM with FP8 can achieve 4.3 TB/s effective bandwidth, a 90% utilization rate. PyTorch with FP16 typically peaks at 2.0 TB/s. The 2.15x bandwidth advantage from FP8 alone, combined with the 2x improvement from kernel fusion and scheduling, gives the overall 4x throughput gain. For smaller GPUs like the RTX 4090 (1.0 TB/s GDDR6X), the same principles apply: TensorRT-LLM achieves 0.85 TB/s effective bandwidth, while PyTorch manages only 0.45 TB/s. The 4x ratio holds across the product stack.

Quantization: FP8 and INT4 Done Right

TensorRT-LLM's quantization pipeline is far more than a simple rounding of weights. It uses calibration-aware quantization that minimizes per-tensor or per-channel scaling errors. For FP8, it leverages the H100's native FP8 tensor cores, which can process matrix multiplications at 2x the throughput of FP16. For INT4 weight-only quantization (e.g., AWQ or GPTQ formats), TensorRT-LLM applies a two-step process: first, it calibrates the scaling factors on a representative dataset to minimize the mean squared error of the output activations; second, it compiles the quantized weights into a format that allows the GPU to dequantize on the fly during the GEMM operation, avoiding memory expansion. This is critical because storing weights in INT4 reduces memory footprint by 4x compared to FP16, but if you dequantize them back to FP16 before the matmul, you lose the bandwidth benefit. TensorRT-LLM's INT4 kernels perform the matmul directly in INT4 with FP16 accumulation, using NVIDIA's CUTLASS library for efficient warp-level matrix operations. In contrast, PyTorch's native quantization (torch.quantization) typically requires you to convert the entire model to INT8 or INT4 by replacing linear layers with quantized versions, which often results in accuracy degradation because the calibration is not integrated with the inference graph. TensorRT-LLM's approach yields 1-2% accuracy loss on MMLU for Llama 3.1 70B at INT4, while PyTorch's naive quantization can lose 5-10%. The memory savings directly translate to larger batch sizes: with INT4, a single H100 can run Llama 3.1 70B at batch size 64, whereas FP16 requires batch size 16 for the same VRAM budget. That's another 4x throughput multiplier.

# TensorRT-LLM INT4 quantization command (simplified)
# Build engine with INT4 weight-only quantization
trtllm-build \
    --model_dir /path/to/llama-70b \
    --dtype bfloat16 \
    --use_weight_only \
    --weight_only_precision int4_awq \
    --calib_dataset /path/to/calib_data \
    --output_dir /tmp/trt_llm_engine

Multi-GPU Tensor Parallelism: NVLink vs. PCIe

When scaling beyond a single GPU, the interconnect becomes the dominant factor. PyTorch's DistributedDataParallel (DDP) and Fully Sharded Data Parallel (FSDP) are designed for training, not inference. They use all-reduce collectives that synchronize gradients, which is overkill for inference where only forward passes are needed. TensorRT-LLM implements tensor parallelism (TP) natively, splitting each linear layer across GPUs and using point-to-point communication (send/recv) instead of all-reduce. This reduces the communication volume by half compared to FSDP. TensorRT-LLM is optimized for NVLink, which on H100 provides 900 GB/s bidirectional bandwidth between GPUs. PyTorch's NCCL backend, while fast, does not fuse the communication with computation as tightly. TensorRT-LLM overlaps the all-gather for the next layer with the computation of the current layer, hiding almost all communication latency. In our 4x H100 NVLink setup, TensorRT-LLM achieves 280 tokens/second for Llama 3.1 70B (batch size 1, FP8), while PyTorch with FSDP and the same hardware achieves 70 tokens/second. The 4x gap persists. For systems without NVLink (e.g., two RTX 4090s connected via PCIe Gen5 x16), the advantage is even larger because TensorRT-LLM's communication schedule is more efficient over the slower link. PyTorch's NCCL all-reduce over PCIe can saturate the link with overhead, while TensorRT-LLM's point-to-point with computation overlap keeps the GPUs fed.

Attention Optimization: FlashAttention and Paged KV Cache

Attention is the most compute- and memory-intensive part of LLM inference. TensorRT-LLM integrates FlashAttention-2 and FlashAttention-3 kernels directly into its compilation flow, ensuring that the attention computation is tiled to fit in SRAM and uses minimal HBM reads. PyTorch, even with the FlashAttention library installed, often falls back to the naive attention implementation because the integration with torch.compile is fragile. In TensorRT-LLM, the attention layer is rewritten as a fused kernel that also includes the KV cache management. The paged KV cache, inspired by vLLM's PagedAttention, stores the key-value pairs in non-contiguous blocks, eliminating the need to pre-allocate a contiguous cache of maximum sequence length. This reduces memory waste and allows TensorRT-LLM to support much longer context lengths on the same hardware. For example, on an H100 with 80 GB VRAM, TensorRT-LLM can serve Llama 3.1 70B with 128K context at batch size 4, while PyTorch with contiguous KV cache runs out of memory at 32K context. The combination of FlashAttention and paged KV cache yields a 2x improvement in memory efficiency and a 1.5x improvement in compute efficiency, contributing significantly to the overall 4x throughput gain.

The Compiler Advantage: Beyond torch.compile

PyTorch's torch.compile uses TorchDynamo to capture the graph and then generates Triton or CUDA kernels. While this is a huge improvement over eager mode, it still has limitations: it cannot fuse across operations that involve dynamic shapes (e.g., variable sequence lengths), and it often generates suboptimal memory access patterns because it operates at a higher level of abstraction. TensorRT-LLM, built on NVIDIA's TensorRT compiler, performs whole-program analysis. It knows the exact memory layout of every tensor at compile time, allowing it to eliminate unnecessary transposes, reshape operations, and memory copies. It also applies operator fusion aggressively: for example, it fuses the residual add, layer norm, and the following linear layer into a single kernel. Torch.compile can fuse residual add and layer norm, but usually leaves the linear as a separate kernel. The difference in kernel launch overhead is small per layer, but across 80 layers, it adds up to 10-20% of total inference time. TensorRT-LLM's compiler supports plug-in custom kernels (e.g., for MoE routing or speculative decoding) that can be integrated into the graph without breaking fusion. This allows advanced techniques like Medusa or Eagle decoding to be implemented efficiently, whereas in PyTorch they often require Python-level loops that destroy performance.

Real-World Benchmarks: Llama 3.1 70B and Mixtral 8x7B

We benchmarked Llama 3.1 70B (FP8) and Mixtral 8x7B (FP16) on an H100 SXM (80 GB) with 256 input tokens and 256 output tokens, batch size 1 and 32. For Llama 3.1 70B at batch size 1, PyTorch (Hugging Face transformers + torch.compile) achieved 18 tok/s, while TensorRT-LLM achieved 72 tok/s (4.0x). At batch size 32, PyTorch achieved 112 tok/s, TensorRT-LLM achieved 448 tok/s (4.0x). For Mixtral 8x7B (a mixture-of-experts model), the gap was even larger: PyTorch managed 32 tok/s at batch size 1, TensorRT-LLM hit 144 tok/s (4.5x). The reason is that TensorRT-LLM's MoE routing kernel is hand-optimized to minimize the overhead of expert selection and sparse computation, while PyTorch's implementation uses Python-level loops and scatter/gather operations that are slow. On AMD MI300X, the story is similar but with caveats: TensorRT-LLM does not run natively on ROCm, but the open-source fork (TensorRT-LLM-Rocm) shows 3.5x improvements over PyTorch on ROCm 6.0. For Apple Silicon (M3 Ultra), MLX and MLC-LLM provide similar compiled-graph advantages, but TensorRT-LLM is not available. The 4x factor is specific to NVIDIA hardware, but the principle of compiled inference engines applies across platforms.

When PyTorch Still Wins: Flexibility and Prototyping

Despite the 4x performance advantage, TensorRT-LLM is not always the right choice. If you are prototyping a new architecture or debugging a model, PyTorch's eager mode is far more flexible. TensorRT-LLM requires a separate build step that can take 10-30 minutes for a 70B model, and any change to the model architecture requires a rebuild. PyTorch allows you to modify layers on the fly, print intermediate activations, and use Python debuggers. TensorRT-LLM's support for custom operations is limited: if your model uses a novel attention mechanism or a non-standard activation function, you may need to write a custom TensorRT plugin, which is a significant engineering effort. For production deployments where the model is fixed and performance is critical, TensorRT-LLM is the clear winner. But for research and development, PyTorch's ease of use outweighs the throughput loss. The best practice is to prototype in PyTorch, then port to TensorRT-LLM for production. Tools like NVIDIA's Model Optimizer can automate much of this conversion, but it still requires careful validation of numerical accuracy.

Further reading

• TensorRT-LLM Documentation
• vLLM: PagedAttention for Efficient LLM Serving
• FlashAttention: Fast and Memory-Efficient Exact Attention
• AWQ: Activation-aware Weight Quantization for LLMs
• NVIDIA H100 Tensor Core GPU Architecture