TensorRT Compiler#

The TensorRT compiler can sweep through the graph to choose the best kernel for each operation and available GPU. Crucially, it can also identify patterns in the graph where multiple operations are good candidates for being fused into a single kernel. This reduces the required amount of memory movement and the overhead of launching multiple GPU kernels.

TensorRT also compiles the graph of operations into a singleCUDA Graph that can be launched all at one time, further reducing the kernel launch overhead.

The TensorRT compiler is extremely powerful for fusing layers and increasing execution speed, but there are some complex layer fusions—likeFlashAttention — that involve interleaving many operations together and which can’t be automatically discovered. For those, you can explicitly replace parts of the graph withplugins at compile time.

Model Engine#

The engine file contains the information that you need for executing the model, but LLM usage in practice requires much more than a single forward pass through the model. TensorRT-LLM includes a highly optimized C++ runtime for executing built LLM engines and managing processes like sampling tokens from the model output, managing the KV cache, and batching requests together.

You can use that runtime directly to execute the model locally, or you can use the TensorRT-LLM runtime backend for NVIDIA Triton Inference Server to serve the model for multiple users.

Weight Bindings#

TensorRT engines embed the network weights, that must be known for compilation.For that reason, the weights must be bound to parameters in the modeldefinition before callingtensorrt_llm.Builder.build_engine. It leads to code like:

# The Linear operator exposes two parameters (see tensorrt_llm/layers/linear.py):classLinear(Module):def__init__(self,...):self.weight=Parameter(shape=(self.out_features,self.in_features),dtype=dtype)self.bias=Parameter(shape=(self.out_features,),dtype=dtype)# The parameters are bound to the weights before compiling the model. See examples/models/core/gpt/weight.py:tensorrt_llm_gpt.layers[i].mlp.fc.weight.value=fromfile(...)tensorrt_llm_gpt.layers[i].mlp.fc.bias.value=fromfile(...)

Note that TensorRT can alsorefitengines to update the weights after compilation. This feature is available toTensorRT-LLM users through therefit_engine method in thetensorrt_llm.Builder class.

Pattern-Matching and Fusion#

One of the key steps performed by TensorRT when it compiles the network graphis the fusion of operations. Fusion is a well-known technique to improve theefficiency when executing LLMs. It helps reduce the amount of data transferredbetween the memory (DRAM) and the compute cores (CUDA cores as well as TensorCores located on theStreamingMultiprocessorsof a GPU). It also removes kernel launch overhead (each time a kernel islaunched on the GPU, there is a small additional CPU cost that is called thelaunch overhead). A classical example is the fusion of the activation functionwith the matrix multiplication (matmul) that usually precedes it in thenetwork.

In TensorRT-LLM, when defining the model, such a sequence can be written as:

c=tensorrt_llm.functional.matmul(a,b)c=tensorrt_llm.functional.relu(c)

During inference, if the above sequence is executed without fusion, thectensor has to be written to global memory at the end of thematmul, read fromthat same memory inrelu and written again afterrelu. If no otheroperation uses the intermediate values betweenmatmul andrelu, it issuboptimal. That is why, during compilation, TensorRT will identify thatpattern and automatically produce a GPU kernel that appliesrelu at the endofmatmul without an intermediate step through global memory. With thatoptimization, thec tensor is written only once (afterrelu) instead oftwice, and is not read between the two operations.

The process of identifying the sequences of operations that can be fused iscalledpattern-matching. TensorRT has a powerful pattern-matching algorithmthat can identify a lot of possible fusions. All the identified patterns areconverted into more efficient kernels by an advanced kernel compiler.

Plugins#

The number of possible fusions is almost infinite and some useful fusionsinvolve very advanced modifications of the graph. A well-known exampleis theFlash-Attention technique tooptimize theMultihead-Attention blockfound in many LLMs. Flash-Attention requires modifications to the arithmeticperformed in the sequenceBMM-Softmax-BMM (whereBMM stands for BatchedMatrix-Matrix product) and the interleaving of thefor-loops of the twobatched matrix products. That’s non-trivial and not necessarily somethingyou can expect a compiler to “discover” on its own (or it might require thesupport for apolyhedralmodel).

As a result, even if TensorRT has a powerful pattern-matching algorithm andsupports a lot of possible fusions, there is always the risk that it cannotidentify uncommon and/or very advanced patterns. To overcome that inevitablelimitation, TensorRT offers a powerful mechanism known asplugins.

The plugins are nodes inserted in the network graph definition that map to user-definedGPU kernels. TensorRT-LLM uses a number of such plugins. They can be found inthecpp/tensorrt_llm/plugins directory.

Plugins are written in C++ and follow a well-defined interface described in theExtending TensorRT with Custom Layerssection of the TensorRTDeveloper Guide.When executed within a TensorRT engine, plugins trigger the execution oftheir encapsulated GPU kernels. A fairly simple example of plugins is theQuantizeTensorPlugin thattriggers a CUDA kernel in theQuantizeTensorPlugin::enqueue member function:

// In cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp:intQuantizeTensorPlugin::enqueue(...){if(inputDesc[0].type==DataType::kFLOAT){invokeQuantization<float>(...);}else{invokeQuantization<half>(...);}return0;}// In cpp/tensorrt_llm/kernels/quantization.cu:template<typenameT>voidinvokeQuantization(...){// The standard <<< >>> construct to launch CUDA kernelsquantizedKernel<<<grid,block,0,stream>>>(...);}

For more details on how TensorRT-LLM implements the GPT Attention operator, seetheMulti-head, Multi-query and Group-query Attention document.

Examples#

Here are examples of Llama 3.1 70B and Llama 3.1 405B showing how to perform multi-GPU and multi-node inference in TensorRT-LLM. The example of Llama 3.1 70B performs multi-GPU inference on a single node, while the example of Llama 3.1 405B performs multi-node inference.

folder_trt_llm=../TensorRT-LLMmodel_dir=Llama-3.1-70Bckpt_dir=ckpt_llama_3.1_70bengine_dir=engine_llama_3.1_70bdtype=bfloat16tp_size=4pp_size=1kv_cache_type=pagedmax_input_len=128max_output_len=128max_batch_size=4workers=$((tp_size*pp_size))python${folder_trt_llm}/examples/models/core/llama/convert_checkpoint.py\--output_dir${ckpt_dir}\--model_dir${model_dir}\--dtype${dtype}\--tp_size${tp_size}\--pp_size${pp_size}\--workers${workers}\--use_parallel_embeddingtrtllm-build\--output_dir${engine_dir}\--checkpoint_dir${ckpt_dir}\--gemm_plugin${dtype}\--gpt_attention_plugin${dtype}\--kv_cache_type${kv_cache_type}\--max_input_len${max_input_len}\--max_seq_len$((max_input_len+max_output_len))\--max_batch_size${max_batch_size}\--workers${workers}

input_text="Born in north-east France, Soyer trained as a"mpirun-n$((tp_size*pp_size))\python${folder_trt_llm}/examples/run.py\--engine_dir${engine_dir}\--tokenizer_dir${model_dir}\--input_text"${input_text}"\--max_output_len${max_output_len}

folder_trt_llm=../TensorRT-LLMmodel_dir=Llama-3.1-405Bckpt_dir=ckpt_llama_3.1_405bengine_dir=engine_llama_3.1_405bdtype=bfloat16tp_size=16pp_size=1kv_cache_type=pagedmax_input_len=128max_output_len=128max_batch_size=4workers=8python${folder_trt_llm}/examples/models/core/llama/convert_checkpoint.py\--output_dir${ckpt_dir}\--model_dir${model_dir}\--dtype${dtype}\--tp_size${tp_size}\--pp_size${pp_size}\--workers${workers}\--use_parallel_embeddingtrtllm-build\--output_dir${engine_dir}\--checkpoint_dir${ckpt_dir}\--gemm_plugin${dtype}\--gpt_attention_plugin${dtype}\--kv_cache_type${kv_cache_type}\--max_input_len${max_input_len}\--max_seq_len$((max_input_len+max_output_len))\--max_batch_size${max_batch_size}\--workers${workers}

#!/bin/bash#SBATCH --account account#SBATCH --partition partition#SBATCH --job-name job-name#SBATCH --time 1:00:00#SBATCH --nodes 2folder_trt_llm=../TensorRT-LLMengine_dir=engine_llama_3.1_405bmodel_dir=Llama-3.1-405Bmax_output_len=128input_text="Born in north-east France, Soyer trained as a"srun\--ntasks-per-node8\--mpipmix\python${folder_trt_llm}/examples/run.py\--engine_dir${engine_dir}\--tokenizer_dir${model_dir}\--input_text"${input_text}"\--max_output_len${max_output_len}

sbatchlaunch_llama_3.1_405b.sh