Attention

Generates an attention.

See also

Multi-Head Attention Fusion

Attributes

normalizationOp activation function can be one of:

  • NONE No normalization is applied.

  • SOFTMAX Apply softmax normalization on the attention scores on the s_kv dimension.

causal (deprecated) boolean parameter determines whether the attention will run causal inference. Superseded by causalKind.

causalKind parameter determines the causal mask alignment orientation. Possible values are:

  • NONE No causal masking applied.

  • UPPER_LEFT Diagonal anchored at top-left corner (legacy default when causal=true).

  • LOWER_RIGHT Diagonal anchored at bottom-right corner (decode-aligned semantics for LLM generation).

decomposable boolean parameter determines whether the attention can be decomposed to use multiple kernels if no fused kernel support found.

normalizationQuantizeToType optional parameter determines the datatype the attention normalization is quantized to.

  • DataType::kFP8: The attention normalization is quantized to FP8.

  • DataType::kINT8: The attention normalization is quantized to INT8.

nbRanks parameter (default: 1) specifies the number of ranks for multi-device attention execution. When > 1, hints attention to perform context-parallel multi-device attention.

Inputs

query a tensor of type T1

key a tensor of type T1

value a tensor of type T1

mask optional: tensor of type T2. For a bool mask, a True value indicates that the corresponding position is allowed to attend. For a float mask, the mask values will be added to the BMM1 output, known as an add mask.

normalizationQuantizeScale optional: tensor of type T1. The quantization scale for the attention normalization output.

Outputs

outputs tensors of type T1

Data Types

T1: float32, float16, bfloat16

T2: float32, float16, bfloat16, bool. T2 must be the same as T1 if not bool

Shape Information

query and outputs are tensors with the same shape of \([b, d_q, s_q, h]\)

key and value are tensors with the same shape of \([b, d_{kv}, s_{kv}, h]\)

mask is a tensor with the same shape of \([a_0, a_1, s_q, s_{kv}]\) where \(a_0\) and \(a_1\) must be broadcastable to b and h.

normalizationQuantizeScale is a tensor with a shape of \([a_0,...,a_n], 0 \leq n \geq 1\)

DLA Support

Not supported.

Multi-Device Runtime Requirements

To use Attention with nbRanks > 1:

  • Multi-device attention requires GPUs with SM >= 100 (Blackwell or newer) and supports only bfloat16 and float16 data types.

  • An NCCL communicator must be initialized and set on the execution context via IExecutionContext::setCommunicator before inference.

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

Examples

Attention
network = get_runner.builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
qkv_shape = (1, 8, 1, 16)
mask_shape = (1, 1, 1, 1)
query = network.add_input("query", dtype=trt.float16, shape=qkv_shape)
key = network.add_input("key", dtype=trt.float16, shape=qkv_shape)
value = network.add_input("value", dtype=trt.float16, shape=qkv_shape)
mask = network.add_input("mask", dtype=trt.bool, shape=mask_shape)
layer = network.add_attention_v2(query, key, value, trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.NONE)
layer.mask = mask
network.mark_output(layer.get_output(0))

query_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    query_data[0, i, 0, :] = i + 1
inputs[query.name] = query_data

key_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    key_data[0, i, 0, :] = i + 1
inputs[key.name] = key_data

value_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    value_data[0, i, 0, :] = i + 1
inputs[value.name] = value_data

mask_data = np.ones(mask_shape, dtype=np.bool_)
inputs[mask.name] = mask_data

outputs[layer.get_output(0).name] = layer.get_output(0).shape

# With identical query/key/values, each position will attend to itself most strongly
expected_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    expected_data[0, i, 0, :] = i + 1
expected[layer.get_output(0).name] = expected_data

# Set get_runner.network back to the new STRONGLY_TYPED network
get_runner.network = network

network = get_runner.builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
qkv_shape = (1, 8, 1, 32)
mask_shape = (1, 1, 1, 1)

query = network.add_input("query", dtype=trt.float16, shape=qkv_shape)
q_scale = network.add_constant(shape=(), weights=np.array(2.0, dtype=np.float16))
query_fp8 = network.add_quantize(query, q_scale.get_output(0), trt.fp8)
query_fp16 = network.add_dequantize(query_fp8.get_output(0), q_scale.get_output(0), trt.float16)

key = network.add_input("key", dtype=trt.float16, shape=qkv_shape)
k_scale = network.add_constant(shape=(), weights=np.array(3.0, dtype=np.float16))
key_fp8 = network.add_quantize(key, k_scale.get_output(0), trt.fp8)
key_fp16 = network.add_dequantize(key_fp8.get_output(0), k_scale.get_output(0), trt.float16)

value = network.add_input("value", dtype=trt.float16, shape=qkv_shape)
v_scale = network.add_constant(shape=(), weights=np.array(4.0, dtype=np.float16))
value_fp8 = network.add_quantize(value, v_scale.get_output(0), trt.fp8)
value_fp16 = network.add_dequantize(value_fp8.get_output(0), v_scale.get_output(0), trt.float16)

mask = network.add_input("mask", dtype=trt.bool, shape=mask_shape)

normalization_quantize_scale = network.add_constant(shape=(), weights=np.array(5.0, dtype=np.float16))
layer = network.add_attention_v2(query_fp16.get_output(0), key_fp16.get_output(0), value_fp16.get_output(0), trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.NONE)

layer.mask = mask
layer.normalization_quantize_scale = normalization_quantize_scale.get_output(0)
layer.normalization_quantize_to_type = trt.fp8

output = layer.get_output(0)
output_scale = network.add_constant(shape=(), weights=np.array(6.0, dtype=np.float16))
output_fp8 = network.add_quantize(output, output_scale.get_output(0), trt.fp8)
output_fp16 = network.add_dequantize(output_fp8.get_output(0), output_scale.get_output(0), trt.float16)
network.mark_output(output_fp16.get_output(0))

query_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    query_data[0, i, 0, :] = i + 1
inputs[query.name] = query_data

key_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    key_data[0, i, 0, :] = i + 1
inputs[key.name] = key_data

value_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    value_data[0, i, 0, :] = i + 1
inputs[value.name] = value_data

mask_data = np.ones(mask_shape, dtype=np.bool_)
inputs[mask.name] = mask_data

outputs[output_fp16.get_output(0).name] = output_fp16.get_output(0).shape

# With identical query/key/values, each position will attend to itself most strongly
expected_data = np.ones(qkv_shape, dtype=np.float16)
for i in range(qkv_shape[1]):
    expected_data[0, i, 0, :] = i + 1
expected[output_fp16.get_output(0).name] = expected_data

# Set get_runner.network back to the new STRONGLY_TYPED network
get_runner.network = network

# Adjust tolerance for FP8 quantization precision loss
get_runner.atol = 0.3

network = get_runner.builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))

# s_q=4, s_kv=8 to test the difference between upper-left and lower-right
batch, heads, s_q, s_kv, head_dim = 1, 1, 4, 8, 16
query_shape = (batch, heads, s_q, head_dim)
kv_shape = (batch, heads, s_kv, head_dim)

query = network.add_input("query", dtype=trt.float16, shape=query_shape)
key = network.add_input("key", dtype=trt.float16, shape=kv_shape)
value = network.add_input("value", dtype=trt.float16, shape=kv_shape)

# Use CausalMaskKind.UPPER_LEFT
layer = network.add_attention_v2(query, key, value, trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.UPPER_LEFT)
layer.decomposable = True
network.mark_output(layer.get_output(0))

# Input data: all ones
query_data = np.ones(query_shape, dtype=np.float16)
key_data = np.ones(kv_shape, dtype=np.float16)
# Value data: each key position has a distinct value
value_data = np.zeros(kv_shape, dtype=np.float16)
for j in range(s_kv):
    value_data[0, 0, j, :] = j + 1  # key pos j has value (j+1)

inputs[query.name] = query_data
inputs[key.name] = key_data
inputs[value.name] = value_data

outputs[layer.get_output(0).name] = layer.get_output(0).shape

# Expected output for upper-left causal:
# Query pos 0: attends to key[0] only -> output = value[0] = 1
# Query pos 1: attends to key[0:1] uniformly -> output = avg(1, 2) = 1.5
# Query pos 2: attends to key[0:2] uniformly -> output = avg(1, 2, 3) = 2
# Query pos 3: attends to key[0:3] uniformly -> output = avg(1, 2, 3, 4) = 2.5
expected_data = np.zeros(query_shape, dtype=np.float16)
for i in range(s_q):
    # Upper-left: attend to positions 0..i
    attended_values = [j + 1 for j in range(i + 1)]
    expected_data[0, 0, i, :] = np.mean(attended_values)

expected[layer.get_output(0).name] = expected_data

get_runner.network = network
get_runner.atol = 0.1

network = get_runner.builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))

# s_q=4, s_kv=8 to test the difference between upper-left and lower-right
batch, heads, s_q, s_kv, head_dim = 1, 1, 4, 8, 16
query_shape = (batch, heads, s_q, head_dim)
kv_shape = (batch, heads, s_kv, head_dim)

query = network.add_input("query", dtype=trt.float16, shape=query_shape)
key = network.add_input("key", dtype=trt.float16, shape=kv_shape)
value = network.add_input("value", dtype=trt.float16, shape=kv_shape)

# Use CausalMaskKind.LOWER_RIGHT
layer = network.add_attention_v2(query, key, value, trt.AttentionNormalizationOp.SOFTMAX, trt.CausalMaskKind.LOWER_RIGHT)
layer.decomposable = True
network.mark_output(layer.get_output(0))

# Input data: all ones for Q and K
query_data = np.ones(query_shape, dtype=np.float16)
key_data = np.ones(kv_shape, dtype=np.float16)
# Value data: each key position has a distinct value
value_data = np.zeros(kv_shape, dtype=np.float16)
for j in range(s_kv):
    value_data[0, 0, j, :] = j + 1  # key pos j has value (j+1)

inputs[query.name] = query_data
inputs[key.name] = key_data
inputs[value.name] = value_data

outputs[layer.get_output(0).name] = layer.get_output(0).shape

# Expected output for lower-right causal:
# offset = s_kv - s_q = 4
# Query pos 0: attends to key[0:4] -> output = avg(1,2,3,4,5) = 3
# Query pos 1: attends to key[0:5] -> output = avg(1,2,3,4,5,6) = 3.5
# Query pos 2: attends to key[0:6] -> output = avg(1,2,3,4,5,6,7) = 4
# Query pos 3: attends to key[0:7] -> output = avg(1,2,3,4,5,6,7,8) = 4.5
offset = s_kv - s_q
expected_data = np.zeros(query_shape, dtype=np.float16)
for i in range(s_q):
    # Lower-right: attend to positions 0..(i + offset)
    attended_values = [j + 1 for j in range(i + offset + 1)]
    expected_data[0, 0, i, :] = np.mean(attended_values)

expected[layer.get_output(0).name] = expected_data

get_runner.network = network
get_runner.atol = 0.1

C++ API

For more information about the C++ IAttention operator, refer to the C++ IAttention documentation.

Python API

For more information about the Python IAttention operator, refer to the Python IAttention documentation.