Working with Transformers

Rotary Position Embedding

TensorRT includes built-in support for RoPE (Rotary Position Embedding) for transformers to make it easier to express RoPE and convert ONNX models with the IRotaryEmbeddingLayer (C++, Python) API to TensorRT.

The IRotaryEmbeddingLayer has three inputs, one optional input, and two attributes:

Inputs

• (index 0) input: The input activation tensor with shape [B, N, S, H]

• (index 1) cosCache: The cosine values for calculating the rotary embedding

• (index 2) sinCache: The sine values for calculating the rotary embedding

cosCache and sinCache should have the shape [B, S, H / 2].

Optional input

• (index 3) positionIds: Position IDs for indexing into cosCache and sinCache

positionIds should have shape [B, S]. When positionIds is provided, cosCache and sinCache should correspondingly have shape [maxPositionId + 1, H / 2].

Attributes

• interleaved: A boolean that specifies whether the input tensor is in interleaved format, that is, whether the 2d vectors rotated are taken from adjacent 2 elements in the hidden dimension.

• rotaryEmbeddingDim: An integer specifying the hidden dimension that participates in RoPE. A special value of 0 means the full hidden dimension participates in RoPE. If it is not 0, then the last dimension of cosCache and sinCache should correspondingly be rotaryEmbeddingDim / 2 instead of H / 2.

The IRotaryEmbeddingLayer has one output, which is the output activation tensor. It has the same shape and format as (index 0) input: the input activation tensor with shape.

The IRotaryEmbeddingLayer is supported by all SM versions.

KV Cache

TensorRT supports using KV cache when doing LLM inference with transformers using the IKVCacheUpdateLayer (C++, Python).

The IKVCacheUpdateLayer has three inputs, one optional input, one output, and two attributes:

Inputs

• (index 0) cache: The K/V cache tensor with shape [B, N, S_max, H]. Allocate this tensor and provide it as an input to the TensorRT network.

• (index 1) update: The newly calculated K or V tensor. Its shape depends on the updateForm attribute:

  • kPADDED_BHND (default): [B, N, S_new, H]

  • kPACKED_NHD: [totalTokens, N, H]

• (index 2) writeIndices: A tensor of shape [B] indicating where to start writing K/V updates for each sequence.

Optional input

• (index 3) updateLengths: Cumulative token counts with shape [B + 1]. The first element must be 0 and the last element equals totalTokens, so the number of tokens for batch i is updateLengths[i + 1] - updateLengths[i]. Required when updateForm is kPACKED_NHD. Use setUpdateLengths to set this tensor.

Attributes

• cacheMode: An enum that specifies the K/V cache update strategy. Currently, only kLINEAR is supported, which performs sequential updates to the cache based on the provided write indices.

• updateForm: An enum (AttentionIOForm) that specifies the layout of the update tensor. The default is kPADDED_BHND. When set to kPACKED_NHD, the update tensor shape is [totalTokens, N, H] and the updateLengths tensor must be provided.

Output

The IKVCacheUpdateLayer has one output: the updated K/V cache tensor with shape [B, N, S_max, H].

Note

IKVCacheUpdateLayer does not own or allocate the cache memory. Memory management is handled by higher-level frameworks (such as TensorRT Edge-LLM) or by your application.

Warning

The cache input tensor and the output tensor must share the same device memory to enable in-place updates. Use IExecutionContext::setTensorAddress to explicitly bind both tensors to the same memory address. If different addresses are assigned, TensorRT reports an API error.

TensorRT also exposes the engine API ICudaEngine::getAliasedInputTensor(char const* outputTensorName) (available in C++ and Python) to retrieve the name of the aliased input tensor when an output tensor is required to be aliased with an input tensor.

The IKVCacheUpdateLayer has the same hardware support matrix as IAttention.

MoE (Mixture of Experts)

TensorRT includes built-in support for MoE (Mixture of Experts) in transformer models using the IMoELayer (C++, Python) API.

The IMoELayer is supported on SM100, SM103, and SM110.

Inputs

• (index 0) hiddenStates: Activation tensor with shape [B, S, H]

• (index 1) selectedExpertsForTokens: Indexes of chosen experts per token, with shape [B, S, topK]

• (index 2) scoresForSelectedExperts: Scores per chosen expert for the weighted output, with shape [B, S, topK]

The following inputs are set using setGatedWeights. Each tensor corresponds to a layer in the expert’s GLU (gated linear unit), where I is the internal hidden size within each expert.

• fcGateWeights: GLU gate weights, with shape [numExperts, H, I]

• fcUpWeights: GLU up-projection weights, with shape [numExperts, H, I]

• fcDownWeights: GLU down-projection weights, with shape [numExperts, I, H]

The following optional inputs are set using setGatedBiases:

• fcGateBias: GLU gate biases, with shape [numExperts, I]

• fcUpBias: GLU up-projection biases, with shape [numExperts, I]

• fcDownBias: GLU down-projection biases, with shape [numExperts, H]

The following input is set using setQuantizationStatic:

• fcDownActivationScale: Quantization scales for the fcDown activation input (mul output), with shape [] or [numExperts]

Attributes

• activationType: Activation type in the expert GLU. Supported values: None and SiLU. Set via setGatedWeights.

• quantizationToType: Type to quantize the mul output (fcDown input) to. Currently only FP8 is supported. Set via setQuantizationStatic.

Note

Additional setters are available for advanced configuration:

• setQuantizationDynamicDblQ: Configures double quantization in the dynamic-quantization path after the mul op. Takes the fcDownActivationDblQScale, dataType, blockShape, and dynQOutputScaleType parameters.

• setSwigluParams(limit, alpha, beta): Configures SwiGLU activation parameters. Defaults: limit = +inf, alpha = 1.0, beta = 0.0.

Output

The IMoELayer produces one output with the same shape as hiddenStates ([B, S, H]), representing the weighted sum of the top-K expert outputs.

Supported Configuration

Parameter	Requirement
Hardware	SM100, SM103, and SM110
hiddenStates	FP16, BF16, FP32, or FP8 static quantization via a DQ layer before the MoE layer.
Weights (unquantized)	FP16, BF16, or FP32
Weights (NVFP4 double quantized)	Weights pass through a DQ layer (with scales from a second DQ), then optionally a Transpose, before the MoE layer. Scales must be FP8. Only available when hiddenStates uses FP8 static quantization.
Internal mul output (fcDown input)	Optional. FP8 static quantization via setQuantizationStatic, or FP4 dynamic double quantization via setQuantizationDynamicDblQ. When not set, no internal quantization is applied.

Fused Attention

TensorRT supports two different methods to trigger Attention fusion:

• Method 1 (recommended): Adding an attention operator to the network using IAttention API.

• Method 2: Constructing an attention graph using primitive INetwork layers like IMatrixMultiplyLayer and ISoftMaxLayer.

Attention computes softmax(Q * Kᵀ ⊗ mask) * V, where:

• Q is query embedding

• K is key embedding

• V is value embeddings

• mask can be:

  • A boolean mask indicating that the corresponding position is allowed to attend by Softmax.

  • An add mask offsetting Q * Kᵀ

Note

TensorRT requires Q or K to be multiplied with \(attention scale\) or both Q and K to be multiplied with \(\sqrt{attention scale}\) before going into IAttention or the first Batched Matrix Multiply (BMM) IMatrixMultiplyLayer for numeric stability and performance reasons.

The shape of Q is [B, N_q, S_q, H], and the shapes of K and V are [B, N_kv, S_kv, H], where:

• B is batch size

• N_q and N_kv are the numbers of attention heads of the query and key/value, respectively

• S_q and S_kv are the sequence lengths of the query and key/value, respectively.

• H is the head/hidden size

Multi-Head Attention (MHA, N_q == N_kv), Group-Query Attention (GQA, N_q % N_kv == 0) and Multi-Query Attention (MQA, N_kv == 1) fusion are all supported.

IO Form (Tensor Layout)

By default, IAttention expects query, key, and value tensors in the padded layout (kPADDED_BHND) with shape [B, N, S, H]. TensorRT also supports a packed layout (kPACKED_NHD) in which all sequences in the batch are concatenated end-to-end without padding, with shape [totalTokens, N, H]. The layout is controlled independently for query and key/value using setQueryForm and setKeyValueForm.

When using the packed layout, you must provide cumulative token counts via setQueryLengths (for the query) and/or setKeyValueLengths (for the key/value). The tensor has shape [B + 1]: the first element must be 0 and the last element equals totalTokens, so the number of tokens for batch i is lengths[i + 1] - lengths[i].

When using the padded layout (kPADDED_BHND), setKeyValueLengths may optionally be used to specify per-batch key/value lengths with a tensor of shape [B]. Each element must not exceed the sequence-length dimension of the KV tensor; if not set, the full sequence-length dimension applies to every batch.

The packed layout removes padding overhead, which is especially beneficial for batches with widely varying sequence lengths.

Supported Attention Fusions

We highly recommend tailoring your model to the restrictions in the following tables, so that Attention fusion happens. It is important because Attention fusion optimizes large sequence length cases by significantly reducing memory footprint from O(S^2) to O(S), where S is the sequence length. On top of that, it shares the common performance benefits of operator fusion, that is, reduced memory traffic, better hardware utilization, less kernel launch, and synchronization overhead.

Table 14 Supported Attention Fusions on SM100, SM103, and SM110

Feature	FP16	BF16	FP8
SM Version (V)	SM100 SM103 SM110	SM100 SM103 SM110	SM100 SM103 SM110
Head Size (H)	8 ≤ H ≤ 128 H % 8 ==0	8 ≤ H ≤ 128 H % 8 ==0	16 ≤ H ≤ 128 H % 16 ==0
Sequence Length (S_q, S_kv)	No restriction	No restriction	No restriction
Quantization	Not required	Not required	Specify Q/DQ layers following the Attention fusion pattern.
IO Form	kPADDED_BHND kPACKED_NHD	kPADDED_BHND kPACKED_NHD	kPADDED_BHND
Accumulation Precision (BMM1)	FP32	FP32	FP32
Accumulation Precision (BMM2)	FP32	FP32	FP32
Supported Mask Type	1d or 2d vector or scalar	1d or 2d vector or scalar	1d or 2d vector or scalar
Pointwise op	Activation, Constant, Elementwise (including SiLU), Scale, and Unary	Activation, Constant, Elementwise (including SiLU), Scale, and Unary	Activation, Constant, Elementwise (including SiLU), Scale, and Unary

Table 15 Supported Attention Fusion for Other SM Versions

Feature	FP16	BF16	INT8	FP8
SM Version (V)	SM75 ≤ V ≤ SM90 SM120 SM121	SM80 ≤ V ≤ SM90 SM120 SM121	SM75 ≤ V ≤ SM90 SM120 SM121	SM89 SM90 SM120 SM121
Head Size (H)	16 ≤ H ≤ 256 H % 8 ==0	16 ≤ H ≤ 256 H % 8 ==0	16 32 64	32 ≤ H ≤ 256 H % 16 == 0
Sequence Length (S_q, S_kv)	No restriction	No restriction	S_{q,kv} ≤ 512	S_{q} = 1 or S_{q,kv} ≥ 32
Quantization	Not required	Not required	Specify Q/DQ layers in the Attention fusion pattern for FP8 and INT8.	Specify Q/DQ layers in the Attention fusion pattern for FP8 and INT8.
IO Form	kPADDED_BHND	kPADDED_BHND	kPADDED_BHND	kPADDED_BHND
Accumulation Precision (BMM1)	FP16 FP32	FP32	INT32	FP32
Accumulation Precision (BMM2)	FP16 FP32 (if ≥ SM90 and S_q = S_kv)	FP32	INT32	FP32
Supported Mask Type	Any masking (such as Select operator in TensorRT)	Any masking (such as Select operator in TensorRT)	Any masking (such as Select operator in TensorRT)	Any masking (such as Select operator in TensorRT)
Pointwise op	Activation, Constant, Elementwise (including SiLU), Pointwise (single input), Scale, and Unary	Activation, Constant, Elementwise (including SiLU), Pointwise (single input), Scale, and Unary	Activation, Constant, Elementwise (including SiLU), Pointwise (single input), Scale, and Unary	Activation, Constant, Elementwise (including SiLU), Pointwise (single input), Scale, and Unary

Note

FP8 fused Attention is expected to outperform FP16 fused Attention when the head size (H) ≥ 128 or the sequence length S_{q,kv} ≥ 128.

Note

Numerical issues in multi-head attention implementations

Attention implementations that use online softmax (for example, Flash Attention) can produce NaN outputs with certain attention masks.

• The attention kernel guards against NaN values at critical computation points.

• Multiply+Add to fused-multiply-add optimization is disabled where unsafe, consistent with other frameworks.

These mitigations may reduce performance by 2–3%, or up to 5–10% in worst cases.

Constructing a Fused Attention

Method 1 (Recommended): Use IAttention API

IAttention offers a direct API to incorporate an attention to the network and trigger the Attention fusion in TensorRT. Refer to the IAttention Operator Documentation for more details.

IAttention by default assumes the attention should always use a fused kernel. If the added attention does not comply with the restrictions listed in the tables above, and thus a fused kernel cannot be utilized, the engine build will raise an error. If users want to allow IAttention to fallback to use non-fused kernels, users can set the IAttention to be decomposable by calling the method IAttention::setDecomposable.

To quantize an IAttention, Q/DQ pairs following Explicit Quantization rules must be used for all query, key and value inputs with another normalization quantize scale and a normalization quantize-to data type supplied to IAttention using IAttention::setNormalizationQuantizeScale and IAttention::setNormalizationQuantizeToType respectively. The normalization quantize scale and the normalization quantize-to data type have to match the data types of the Q/DQ pairs. The output of the IAttention can be quantized with an optional quantizer.

To apply masks with IAttention, use IAttention::setCausalKind with CausalMaskKind::kNONE, CausalMaskKind::kUPPER_LEFT, or CausalMaskKind::kLOWER_RIGHT for implicit causal masking, or IAttention::setMask for arbitrary masks (boolean or float), the legacy IAttention::setCausal(bool) helper is deprecated and now internally maps true to kUPPER_LEFT and false to kNONE.

To add a Sage Attention and trigger the fusion, follow the paradigm in the above figure and use IDynamicQuantizeLayer in place of IQuantizeLayer after query, key, value.

Method 2: Construct an Attention Graph with Primitive ``INetwork`` Layers

The following figures demonstrate how FP16, BF16, FP8, and INT8 attentions can be constructed to trigger Attention fusion. Batched MatMul uses IMatrixMultiplyLayer, Softmax uses ISoftMaxLayer and Q/DQ uses IQuantizeLayer and IDequantizeLayer.

TensorRT chooses the accumulation precision by default based on the input types and performance considerations. To enforce a specific accumulation precision in a strongly typed network, cast the inputs to the desired type using ICastLayer — TensorRT recognizes this pattern and fuses the casts with the operation.

The Attention fusion captures common pointwise operators in series in Attention as mentioned in the pointwise operation list. It also covers Q/DQ fusion following Attention for certain quantization and architecture (such as FP16/BF16 to FP8/INT8 on NVIDIA Ampere GPU architecture).

For method 2, TensorRT can decide not to fuse an Attention graph into a single kernel based on performance evaluations or other constraints.

Example Workflow: Quantize an ONNX Attention Model to FP8 with ModelOpt

Assume you have an ONNX model, vit_base_patch8_224_Opset17.onnx, and calibration data, calib.npy, on your local machine.

1. Install the TensorRT model optimizer.

   pip3 install --no-cache-dir --extra-index-url https://pypi.nvidia.com nvidia-modelop
   

2. Quantize a model with TensorRT model optimizer. For more information, refer to these detailed instructions.

   python3 -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch8_224_Opset17.onnx \
    --quantize_mode=<fp8|int8> \
    --calibration_data=calib.npy \
    --calibration_method=<max|entropy> \
    --output_path=vit_base_patch8_224_Opset17.quant.onnx
   

   For example, to quantize the model using FP8 and entropy and store as vit_base_patch8_224_Opset17.quant.onnx, use:

   python3 -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch8_224_Opset17.onnx \
    --quantize_mode=fp8 \
    --calibration_data=calib.npy \
    --calibration_method=entropy \
    --output_path=vit_base_patch8_224_Opset17.quant.onnx
   

3. Compile the quantized model with TensorRT.

   trtexec --onnx=vit_base_patch8_224_Opset17.quant.onnx \
   --saveEngine=vit_base_patch8_224_Opset17.engine \
   --stronglyTyped --skipInference --profilingVerbosity=detailed
   

4. Run the quantized model with TensorRT.

   trtexec --loadEngine=vit_base_patch8_224_Opset17.engine \
   --useCudaGraph --noDataTransfers --useSpinWait
   

   Add the following options if you want to check if Attention is fused. Attention should be fused if you find the mha op in the output.log file.

   trtexec --loadEngine=vit_base_patch8_224_Opset17.engine \
   --profilingVerbosity=detailed --dumpLayerInfo --skipInference &> output.log
   

   Tip

   There are two ways to set the accumulation data type to FP32:

   1. Manually insert ICastLayer to cast the attention inputs to FP32 before the GEMM. In a strongly typed network, TensorRT fuses the casts with the GEMM and produces a single kernel with FP16 inputs and FP32 accumulation.

   2. Convert your ONNX model using TensorRT Model Optimizer, which adds the Cast ops automatically.

      • If the Attention has a head size (H) that is not a multiple of 16, for better performance do not add Q/DQ ops in the Attention to fall back to the FP16 or BF16 Attention.

      • Compare Attention fusion performance of INT8 with FP8.

Example: Exploiting Sparsity in Attention Masks

In limited cases, TensorRT can exploit block sparsity in attention masks to reduce the runtime of the attention computation. For example, an attention with a causal mask requires half the computation of a regular attention (since it only considers query tokens with index position in the sequence at or after a corresponding key token’s index, for example, q >= kv), therefore, TensorRT can achieve up to 2x speedup on the attention computation.

In ONNX, such as here is how you would express a causal mask for primitive op constructed attention.

While TensorRT’s overall approach to recognizing sparsity in masks is flexible, there are a few characteristics of this expression that must be present for TensorRT to consider such an optimization.

• The mask is expressed using indices along the sequence length (that is, through the Range op). A causal mask is written as q >= kv, therefore, we get the index, the sequence length, the query, key dimensions through the Range ops, and reshape them so they can be compared.

• The mask subgraph (that is, the subgraph of operations feeding into the Add, starting with the Where) must be expressed solely using pointwise operations, position expressions (such as Trilu, Range), or Expand operations (or equivalent Reshape, Tile operations).

• Along the chain of operations from the MatMul to the Softmax, there must be exactly one operation doing masking, in this case the Add. That operation with input dependent on the MatMul vIn and output vOut must obey the property that at a point (q,kv) in the attention computation:

  • If the mask is valid at that point, vOut = vIn.

  • If the mask is invalid at that point, vOut = -inf.

• The masking operation must be a Where, Add, or Subtract operation.

Example: Using Transformer-Oriented APIs

The following examples demonstrate common attention layout patterns using IAttention (optionally combined with IKVCacheUpdateLayer and IRotaryEmbeddingLayer).

Padded Format

C++

ITensor* qInput = network->addInput("q", DataType::kHALF,
                         Dims4{-1, numQHeads, -1, headSize});
ITensor* kInput = network->addInput("k", DataType::kHALF,
                                    Dims4{-1, numKvHeads, -1, headSize});
ITensor* vInput = network->addInput("v", DataType::kHALF,
                                    Dims4{-1, numKvHeads, -1, headSize});

ITensor* kCacheInput = network->addInput("kCache", DataType::kHALF,
                                          Dims4{-1, numKvHeads, kvCacheCapacity, headSize});
ITensor* vCacheInput = network->addInput("vCache", DataType::kHALF,
                                          Dims4{-1, numKvHeads, kvCacheCapacity, headSize});
ITensor* ropeCosCache = network->addInput("ropeCosCache", DataType::kHALF,
                                          Dims2{maxPositionEmbeddings, headSize / 2});
ITensor* ropeSinCache = network->addInput("ropeSinCache", DataType::kHALF,
                                          Dims2{maxPositionEmbeddings, headSize / 2});
ITensor* positionIds = network->addInput("positionIds", DataType::kINT32,
                                          Dims2{-1, -1});
ITensor* kCacheIndices = network->addInput("kCacheIndices", DataType::kINT32,
                                            Dims{1, {-1}});
// Per-batch KV lengths: [B], each element is the number of valid KV tokens for that batch element
ITensor* kvLengths = network->addInput("kvLengths", DataType::kINT32,
                                        Dims{1, {-1}});

// Apply RoPE to Q and K
IRotaryEmbeddingLayer* ropeQ = network->addRotaryEmbedding(*qInput, *ropeCosCache, *ropeSinCache,
                                                            false, 0);
ropeQ->setInput(3, *positionIds);
IRotaryEmbeddingLayer* ropeK = network->addRotaryEmbedding(*kInput, *ropeCosCache, *ropeSinCache,
                                                            false, 0);
ropeK->setInput(3, *positionIds);

ITensor* ropeQTensor = ropeQ->getOutput(0);
ITensor* ropeKTensor = ropeK->getOutput(0);

// Update KV cache
KVCacheMode kvCacheMode = KVCacheMode::kLINEAR;
IKVCacheUpdateLayer* kCacheLayer = network->addKVCacheUpdate(*kCacheInput, *ropeKTensor,
                                                              *kCacheIndices, kvCacheMode);
IKVCacheUpdateLayer* vCacheLayer = network->addKVCacheUpdate(*vCacheInput, *vInput,
                                                              *kCacheIndices, kvCacheMode);

ITensor* kCacheOutput = kCacheLayer->getOutput(0);
ITensor* vCacheOutput = vCacheLayer->getOutput(0);

// Mark cache outputs
kCacheOutput->setName("kCacheOutput");
network->markOutput(*kCacheOutput);
vCacheOutput->setName("vCacheOutput");
network->markOutput(*vCacheOutput);

// Scale Q
float qkScale = 1.0f / std::sqrt(static_cast<float>(headSize));
half_float::half scaleData{qkScale};
IConstantLayer* scaleConst = network->addConstant(Dims4{1, 1, 1, 1},
                                                  Weights{DataType::kHALF, &scaleData, 1});

IElementWiseLayer* scaledQ = network->addElementWise(*ropeQTensor, *scaleConst->getOutput(0),
                                                      ElementWiseOperation::kPROD);
ITensor* scaledQTensor = scaledQ->getOutput(0);

// Add attention layer using the full cache as KV input.
// Use setKeyValueLengths to specify per-batch valid KV lengths instead of
// slicing the cache to extract the present KV.
IAttention* attention = network->addAttentionV2(*scaledQTensor, *kCacheOutput, *vCacheOutput,
                                                AttentionNormalizationOp::kSOFTMAX,
                                                CausalMaskKind::kNONE);
attention->setKeyValueLengths(kvLengths);
attention->setDecomposable(true);

ITensor* attentionOutput = attention->getOutput(0);
attentionOutput->setName("attentionOutput");
network->markOutput(*attentionOutput);

Python

q_input = network.add_input("q", trt.float16,
               (-1, num_q_heads, -1, head_size))
k_input = network.add_input("k", trt.float16,
                          (-1, num_kv_heads, -1, head_size))
v_input = network.add_input("v", trt.float16,
                          (-1, num_kv_heads, -1, head_size))

k_cache_input = network.add_input("k_cache", trt.float16,
                                (-1, num_kv_heads, kv_cache_capacity, head_size))
v_cache_input = network.add_input("v_cache", trt.float16,
                                (-1, num_kv_heads, kv_cache_capacity, head_size))

rope_cos_cache = network.add_input("rope_cos_cache", trt.float16,
                                  (max_position_embeddings, head_size // 2))
rope_sin_cache = network.add_input("rope_sin_cache", trt.float16,
                                  (max_position_embeddings, head_size // 2))

position_ids = network.add_input("position_ids", trt.int32, (-1, -1))
kv_cache_indices = network.add_input("kv_cache_indices", trt.int32, (-1,))
# Per-batch KV lengths: [B], each element is the number of valid KV tokens for that batch element
kv_lengths = network.add_input("kv_lengths", trt.int32, (-1,))

# Apply RoPE to Q and K
rope_q = network.add_rotary_embedding(
    q_input, rope_cos_cache, rope_sin_cache, False, 0)
rope_q.set_input(3, position_ids)

rope_k = network.add_rotary_embedding(
    k_input, rope_cos_cache, rope_sin_cache, False, 0)
rope_k.set_input(3, position_ids)

rope_q_tensor = rope_q.get_output(0)
rope_k_tensor = rope_k.get_output(0)

# Update KV cache
kv_cache_mode = trt.KVCacheMode.LINEAR
k_cache_layer = network.add_kv_cache_update(k_cache_input,
                                            rope_k_tensor,
                                            kv_cache_indices,
                                            kv_cache_mode)
v_cache_layer = network.add_kv_cache_update(v_cache_input,
                                            v_input,
                                            kv_cache_indices,
                                            kv_cache_mode)

k_cache_output = k_cache_layer.get_output(0)
v_cache_output = v_cache_layer.get_output(0)

# Mark cache outputs
k_cache_output.name = "k_cache_output"
network.mark_output(k_cache_output)
v_cache_output.name = "v_cache_output"
network.mark_output(v_cache_output)

# Scale Q
qk_scale = 1.0 / (head_size**0.5)
scale_const = network.add_constant((1, 1, 1, 1),
                                  np.array([qk_scale],
                                  dtype=np.float16))
scaled_q = network.add_elementwise(rope_q_tensor,
                                  scale_const.get_output(0),
                                  trt.ElementWiseOperation.PROD)
scaled_q_tensor = scaled_q.get_output(0)

# Add attention layer using the full cache as KV input.
# Use set_key_value_lengths to specify per-batch valid KV lengths instead of
# slicing the cache to extract the present KV.
attention = network.add_attention_v2(scaled_q_tensor, k_cache_output, v_cache_output,
                                    trt.AttentionNormalizationOp.SOFTMAX,
                                    trt.CausalMaskKind.NONE)
attention.set_key_value_lengths(kv_lengths)
attention.decomposable = True

attention_output = attention.get_output(0)
attention_output.name = "attention_output"
network.mark_output(attention_output)

Packed (Ragged) Format

The following example shows how to use the packed (kPACKED_NHD) format with IAttention and IKVCacheUpdateLayer. In this format, all sequences in the batch are concatenated end-to-end without padding. This example pairs a packed query with a padded KV cache, the typical LLM decode pattern.

C++

// Total tokens across all sequences in the batch
ITensor* qInput = network->addInput("q", DataType::kHALF,
                                     Dims3{-1, numQHeads, headSize});  // [totalTokens, N_q, H]
ITensor* kInput = network->addInput("k", DataType::kHALF,
                                     Dims3{-1, numKvHeads, headSize});  // [totalTokens, N_kv, H]
ITensor* vInput = network->addInput("v", DataType::kHALF,
                                     Dims3{-1, numKvHeads, headSize});  // [totalTokens, N_kv, H]

ITensor* kCacheInput = network->addInput("kCache", DataType::kHALF,
                                          Dims4{-1, numKvHeads, kvCacheCapacity, headSize});
ITensor* vCacheInput = network->addInput("vCache", DataType::kHALF,
                                          Dims4{-1, numKvHeads, kvCacheCapacity, headSize});
ITensor* ropeCosCache = network->addInput("ropeCosCache", DataType::kHALF,
                                           Dims2{maxPositionEmbeddings, headSize / 2});
ITensor* ropeSinCache = network->addInput("ropeSinCache", DataType::kHALF,
                                           Dims2{maxPositionEmbeddings, headSize / 2});
ITensor* positionIds = network->addInput("positionIds", DataType::kINT32,
                                          Dims{1, {-1}});  // [totalTokens]
ITensor* kCacheIndices = network->addInput("kCacheIndices", DataType::kINT32,
                                            Dims{1, {-1}});

// Cumulative token counts: [B + 1], first element is 0, last element is totalTokens
ITensor* queryLengths = network->addInput("queryLengths", DataType::kINT32,
                                           Dims{1, {-1}});
// Per-batch KV lengths: [B], each element is the number of valid KV tokens for that batch element
ITensor* kvLengths = network->addInput("kvLengths", DataType::kINT32,
                                        Dims{1, {-1}});

// Apply RoPE to Q and K (RoPE works with both padded and packed inputs)
IRotaryEmbeddingLayer* ropeQ = network->addRotaryEmbedding(*qInput, *ropeCosCache, *ropeSinCache,
                                                            false, 0);
ropeQ->setInput(3, *positionIds);
IRotaryEmbeddingLayer* ropeK = network->addRotaryEmbedding(*kInput, *ropeCosCache, *ropeSinCache,
                                                            false, 0);
ropeK->setInput(3, *positionIds);

ITensor* ropeQTensor = ropeQ->getOutput(0);
ITensor* ropeKTensor = ropeK->getOutput(0);

// Update KV cache with packed update form
KVCacheMode kvCacheMode = KVCacheMode::kLINEAR;
IKVCacheUpdateLayer* kCacheLayer = network->addKVCacheUpdate(*kCacheInput, *ropeKTensor,
                                                              *kCacheIndices, kvCacheMode);
kCacheLayer->setUpdateForm(AttentionIOForm::kPACKED_NHD);
kCacheLayer->setUpdateLengths(queryLengths);

IKVCacheUpdateLayer* vCacheLayer = network->addKVCacheUpdate(*vCacheInput, *vInput,
                                                              *kCacheIndices, kvCacheMode);
vCacheLayer->setUpdateForm(AttentionIOForm::kPACKED_NHD);
vCacheLayer->setUpdateLengths(queryLengths);

ITensor* kCacheOutput = kCacheLayer->getOutput(0);
ITensor* vCacheOutput = vCacheLayer->getOutput(0);

kCacheOutput->setName("kCacheOutput");
network->markOutput(*kCacheOutput);
vCacheOutput->setName("vCacheOutput");
network->markOutput(*vCacheOutput);

// Scale Q
float qkScale = 1.0f / std::sqrt(static_cast<float>(headSize));
half_float::half scaleData{qkScale};
IConstantLayer* scaleConst = network->addConstant(Dims3{1, 1, 1},
                                                   Weights{DataType::kHALF, &scaleData, 1});

IElementWiseLayer* scaledQ = network->addElementWise(*ropeQTensor, *scaleConst->getOutput(0),
                                                      ElementWiseOperation::kPROD);
ITensor* scaledQTensor = scaledQ->getOutput(0);

// Add attention layer with packed query and padded KV (full cache).
// Use setKeyValueLengths to specify per-batch valid KV lengths instead of
// slicing the cache to extract the present KV.
IAttention* attention = network->addAttentionV2(*scaledQTensor, *kCacheOutput, *vCacheOutput,
                                                 AttentionNormalizationOp::kSOFTMAX,
                                                 CausalMaskKind::kNONE);
attention->setQueryForm(AttentionIOForm::kPACKED_NHD);
attention->setQueryLengths(queryLengths);
attention->setKeyValueLengths(kvLengths);
attention->setDecomposable(true);

ITensor* attentionOutput = attention->getOutput(0);
attentionOutput->setName("attentionOutput");
network->markOutput(*attentionOutput);

Python

q_input = network.add_input("q", trt.float16,
                            (-1, num_q_heads, head_size))  # [totalTokens, N_q, H]
k_input = network.add_input("k", trt.float16,
                            (-1, num_kv_heads, head_size))  # [totalTokens, N_kv, H]
v_input = network.add_input("v", trt.float16,
                            (-1, num_kv_heads, head_size))  # [totalTokens, N_kv, H]

k_cache_input = network.add_input("k_cache", trt.float16,
                                  (-1, num_kv_heads, kv_cache_capacity, head_size))
v_cache_input = network.add_input("v_cache", trt.float16,
                                  (-1, num_kv_heads, kv_cache_capacity, head_size))

rope_cos_cache = network.add_input("rope_cos_cache", trt.float16,
                                   (max_position_embeddings, head_size // 2))
rope_sin_cache = network.add_input("rope_sin_cache", trt.float16,
                                   (max_position_embeddings, head_size // 2))

position_ids = network.add_input("position_ids", trt.int32, (-1,))  # [totalTokens]
kv_cache_indices = network.add_input("kv_cache_indices", trt.int32, (-1,))

# Cumulative token counts: [B + 1], first element is 0, last element is totalTokens
query_lengths = network.add_input("query_lengths", trt.int32, (-1,))
# Per-batch KV lengths: [B], each element is the number of valid KV tokens for that batch element
kv_lengths = network.add_input("kv_lengths", trt.int32, (-1,))

# Apply RoPE to Q and K (RoPE works with both padded and packed inputs)
rope_q = network.add_rotary_embedding(
    q_input, rope_cos_cache, rope_sin_cache, False, 0)
rope_q.set_input(3, position_ids)

rope_k = network.add_rotary_embedding(
    k_input, rope_cos_cache, rope_sin_cache, False, 0)
rope_k.set_input(3, position_ids)

rope_q_tensor = rope_q.get_output(0)
rope_k_tensor = rope_k.get_output(0)

# Update KV cache with packed update form
kv_cache_mode = trt.KVCacheMode.LINEAR
k_cache_layer = network.add_kv_cache_update(k_cache_input,
                                            rope_k_tensor,
                                            kv_cache_indices,
                                            kv_cache_mode)
k_cache_layer.update_form = trt.AttentionIOForm.PACKED_NHD
k_cache_layer.set_update_lengths(query_lengths)

v_cache_layer = network.add_kv_cache_update(v_cache_input,
                                            v_input,
                                            kv_cache_indices,
                                            kv_cache_mode)
v_cache_layer.update_form = trt.AttentionIOForm.PACKED_NHD
v_cache_layer.set_update_lengths(query_lengths)

k_cache_output = k_cache_layer.get_output(0)
v_cache_output = v_cache_layer.get_output(0)

k_cache_output.name = "k_cache_output"
network.mark_output(k_cache_output)
v_cache_output.name = "v_cache_output"
network.mark_output(v_cache_output)

# Scale Q
qk_scale = 1.0 / (head_size ** 0.5)
scale_const = network.add_constant((1, 1, 1),
                                   np.array([qk_scale],
                                            dtype=np.float16))
scaled_q = network.add_elementwise(rope_q_tensor,
                                   scale_const.get_output(0),
                                   trt.ElementWiseOperation.PROD)
scaled_q_tensor = scaled_q.get_output(0)

# Add attention layer with packed query and padded KV (full cache).
# Use set_key_value_lengths to specify per-batch valid KV lengths instead of
# slicing the cache to extract the present KV.
attention = network.add_attention_v2(scaled_q_tensor, k_cache_output, v_cache_output,
                                     trt.AttentionNormalizationOp.SOFTMAX,
                                     trt.CausalMaskKind.NONE)
attention.query_form = trt.AttentionIOForm.PACKED_NHD
attention.set_query_lengths(query_lengths)
attention.set_key_value_lengths(kv_lengths)
attention.decomposable = True

attention_output = attention.get_output(0)
attention_output.name = "attention_output"
network.mark_output(attention_output)

Fully Packed (Ragged) Format

This example uses the packed (kPACKED_NHD) format for all three of query, key, and value. This pattern is useful when there is no KV cache — for example, in vision transformers where each image contributes a different number of tokens to the batch.

C++

// All inputs are packed: [totalTokens, numHeads, headSize]
ITensor* qInput = network->addInput("q", DataType::kHALF,
                                     Dims3{-1, numQHeads, headSize});   // [totalTokens, N_q, H]
ITensor* kInput = network->addInput("k", DataType::kHALF,
                                     Dims3{-1, numKvHeads, headSize});  // [totalTokens, N_kv, H]
ITensor* vInput = network->addInput("v", DataType::kHALF,
                                     Dims3{-1, numKvHeads, headSize});  // [totalTokens, N_kv, H]

// Cumulative token counts: [B + 1], first element is 0, last element is totalTokens
// For a batch of images with token counts [196, 256, 128]:
//   queryLengths = kvLengths = [0, 196, 452, 580]
ITensor* queryLengths = network->addInput("queryLengths", DataType::kINT32,
                                           Dims{1, {-1}});
ITensor* kvLengths = network->addInput("kvLengths", DataType::kINT32,
                                        Dims{1, {-1}});

// Scale Q
float qkScale = 1.0f / std::sqrt(static_cast<float>(headSize));
half_float::half scaleData{qkScale};
IConstantLayer* scaleConst = network->addConstant(Dims3{1, 1, 1},
                                                   Weights{DataType::kHALF, &scaleData, 1});

IElementWiseLayer* scaledQ = network->addElementWise(*qInput, *scaleConst->getOutput(0),
                                                      ElementWiseOperation::kPROD);
ITensor* scaledQTensor = scaledQ->getOutput(0);

// Add attention with fully packed Q, K, V
IAttention* attention = network->addAttentionV2(*scaledQTensor, *kInput, *vInput,
                                                 AttentionNormalizationOp::kSOFTMAX,
                                                 CausalMaskKind::kNONE);
attention->setQueryForm(AttentionIOForm::kPACKED_NHD);
attention->setKeyValueForm(AttentionIOForm::kPACKED_NHD);
attention->setQueryLengths(queryLengths);
attention->setKeyValueLengths(kvLengths);
attention->setDecomposable(true);

ITensor* attentionOutput = attention->getOutput(0);
attentionOutput->setName("attentionOutput");
network->markOutput(*attentionOutput);

Python

# All inputs are packed: [totalTokens, numHeads, headSize]
q_input = network.add_input("q", trt.float16,
                            (-1, num_q_heads, head_size))    # [totalTokens, N_q, H]
k_input = network.add_input("k", trt.float16,
                            (-1, num_kv_heads, head_size))   # [totalTokens, N_kv, H]
v_input = network.add_input("v", trt.float16,
                            (-1, num_kv_heads, head_size))   # [totalTokens, N_kv, H]

# Cumulative token counts: [B + 1], first element is 0, last element is totalTokens
# For a batch of images with token counts [196, 256, 128]:
#   query_lengths = kv_lengths = [0, 196, 452, 580]
query_lengths = network.add_input("query_lengths", trt.int32, (-1,))
kv_lengths = network.add_input("kv_lengths", trt.int32, (-1,))

# Scale Q
qk_scale = 1.0 / (head_size ** 0.5)
scale_const = network.add_constant((1, 1, 1),
                                   np.array([qk_scale],
                                            dtype=np.float16))
scaled_q = network.add_elementwise(q_input,
                                   scale_const.get_output(0),
                                   trt.ElementWiseOperation.PROD)
scaled_q_tensor = scaled_q.get_output(0)

# Add attention with fully packed Q, K, V
attention = network.add_attention_v2(scaled_q_tensor, k_input, v_input,
                                     trt.AttentionNormalizationOp.SOFTMAX,
                                     trt.CausalMaskKind.NONE)
attention.query_form = trt.AttentionIOForm.PACKED_NHD
attention.key_value_form = trt.AttentionIOForm.PACKED_NHD
attention.set_query_lengths(query_lengths)
attention.set_key_value_lengths(kv_lengths)
attention.decomposable = True

attention_output = attention.get_output(0)
attention_output.name = "attention_output"
network.mark_output(attention_output)

Multi-Device Attention

TensorRT supports multi-device attention through context parallelism, splitting the key-value sequence across multiple GPUs. Each GPU processes a portion of the sequence, and NCCL handles inter-device communication. Refer to the attention-mdtrt sample as an example.

• Supported data types: BF16 and FP16

• Platform support: Blackwell (SM 100) and later

For more information, refer to the GPU Architecture and Precision Support section.

Note

There is no ONNX compliant operator for multi-device attention. However, standard ONNX files generated by a sharding tool are supported when the Attention op includes an nbRanks field. You can also use the C++ or Python API directly.

The IAttention layer runs in multi-device mode when nbRanks > 1. Use IAttention::setNbRanks (C++) or IAttention.num_ranks (Python) to set the number of devices. Once set to a value greater than 1, this value cannot be changed.

Build-time configuration

1. Add the attention layer with the IAttention API.

2. Call IAttention::setNbRanks(nbRanks) (C++) or set attention.num_ranks = nbRanks (Python).

Runtime requirements

1. Initialize an NCCL communicator with ncclCommInitRank or ncclCommInitAll.

2. Call IExecutionContext::setCommunicator(ncclComm_t) before inference.

3. All participating ranks must execute the same network with synchronized execution calls.

C++

// Add attention layer and set number of ranks
auto attention = network->addAttentionV2(query, key, value, AttentionNormalizationOp::kSOFTMAX, CausalMaskKind::kNONE);

// At runtime, on each rank, set communicator and run inference
context->setCommunicator(ncclComm);
context->enqueueV3(stream);

Python

# Add attention layer and set number of ranks
attention = network.add_attention_v2(query, key, value, trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.NONE)
attention.num_ranks = num_gpus

# At runtime, on each rank, set communicator and run inference
context.set_communicator(nccl_comm)
context.execute_async_v3(stream)

See also

Working with Loops

Loop constructs for autoregressive decoding patterns.

Working with Conditionals

Conditional execution for dynamic decoder logic.