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 13 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 14 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.

Primitive Layers MHA Fusion Pattern for FP16/BF16

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.

MHA fusion pattern for FP8 and INT8 quantization

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.

Install the TensorRT model optimizer.

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

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

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

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.

An ONNX Graph with Sparse Causal Mask and 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)