TensorRT Custom Operator Schemas

TRT_QuantizedAttention

Custom op for quantized Attention based on ONNX Attention-23. The common attributes, inputs, and outputs between this op and the standard Attention op follow the same rules as in the standard spec; only the ones unique to this custom op are followed by a description. This op supports the same padded and packed/ragged tensor forms as TRT_Attention.

Attributes

query_form - STRING

"padded_bhnd" (default) or "packed_nhd". Sets the IO form of the Q tensor and output. Controls Q independently of KV.

kv_form - STRING

"padded_bhnd" (default) or "packed_nhd". Sets the IO form of the K and V tensors. Controls KV independently of Q.

is_causal

kv_num_heads

q_num_heads

scale

normalization_quantize_to_type - INT

fp8 (default) or int8. The quantization type for Q/DQ between normalization and BMM2.

TRT_decomposable - INT

Default is 0.

TRT_normalization_op - STRING

"softmax" (default) or "none".

See the TRT_Attention tensor format notation for the meaning of "padded_bhnd" and "packed_nhd".

Inputs (5-7)

Q (index 0)

K (index 1)

V (index 2)

attn_mask (optional, index 3)

normalization_quantize_scale (optional, index 4)

Scalar tensor of fp32, fp16, or bf16 type. The quantization scale for Q/DQ between normalization and BMM2.

query_lengths (optional, index 5) - INT32

1D tensor of shape [batchSize + 1] containing cumulative token counts for the query sequence, starting with 0, e.g. [0, len_0, len_0 + len_1, ...]. Must be provided when query_form is "packed_nhd".

kv_lengths (optional, index 6) - INT32

1D tensor whose shape depends on kv_form: when "packed_nhd", shape is [batchSize + 1] containing cumulative token counts starting with 0, e.g. [0, len_0, len_0 + len_1, ...] and this input must be provided; when "padded_bhnd", shape is [batchSize] containing per-batch sequence lengths and this input is optional, e.g. [len_0, len_1, ...].

Outputs

Y

TRT_Attention

Custom op for attention with packed (ragged) batching support. This op mirrors TensorRT's IAttention API directly: Q, K, and V are passed to TensorRT without implicit reshaping or scaling. The caller must provide tensors in the shape required by the selected IO forms.

Tensor format notation: "padded_bhnd" denotes 4D tensors with shape [B, H, N, D] and "packed_nhd" denotes 3D tensors with shape [T, H, D], where B=batch, H=num_heads, N=sequence_length, D=head_size, and T=total_tokens (sum of all sequence lengths in the batch). This notation follows the ONNX Attention-23 spec; the TRT C++ API docs use [b, d, s, h] for the same layout (d=num_heads, s=seq_len, h=head_size).

Attributes

query_form - STRING

"padded_bhnd" (default) or "packed_nhd". Sets the IO form of the Q tensor and output. Controls Q independently of KV.

kv_form - STRING

"padded_bhnd" (default) or "packed_nhd". Sets the IO form of the K and V tensors. Controls KV independently of Q.

causal_kind - STRING

"none" (default), "upper_left", or "lower_right". Sets the implicit causal mask kind.

TRT_decomposable - INT

0 (default) or 1. When set to 1, TensorRT may decompose the attention into multiple kernels if no fused kernel is supported.

TRT_normalization_op - STRING

"softmax" (default) or "none".

nb_rank - INT

Number of ranks for multi-device attention execution. Default is 1.

The query_form and kv_form attributes are independent: Q and KV may use different forms (e.g., packed Q with padded KV for LLM context-phase attention).

Inputs (3-6)

Q (index 0)

K (index 1)

V (index 2)

attn_mask (optional, index 3)

Mask tensor with rank at most 4. A BOOL mask uses true values for positions that may attend; other mask values are added to the BMM1 output.

query_lengths (optional, index 4) - INT32

1D tensor of shape [batchSize + 1] containing cumulative token counts for the query sequence, starting with 0, e.g. [0, len_0, len_0 + len_1, ...]. Must be provided when query_form is "packed_nhd".

kv_lengths (optional, index 5) - INT32

1D tensor whose shape depends on kv_form: when "packed_nhd", shape is [batchSize + 1] containing cumulative token counts starting with 0, e.g. [0, len_0, len_0 + len_1, ...] and this input must be provided; when "padded_bhnd", shape is [batchSize] containing per-batch sequence lengths and this input is optional, e.g. [len_0, len_1, ...].

Outputs

Y

Output tensor with the same shape, IO form, and data type as Q.

Restrictions

• Q, K, and V must already match the selected IO forms. Unlike ONNX Attention, TRT_Attention does not use q_num_heads, kv_num_heads, or scale to reshape or scale Q/K/V.
• attn_mask cannot be used together with causal_kind values other than "none".
• When query and KV forms differ, Q and KV tensors may have different ranks. For example, with packed Q and padded KV: Q is [T, numHeads, headSize] while K/V are [batchSize, numHeads, seqLen, headSize].

TRT_KVCacheUpdate

Custom op for KV cache updates with packed update support. Based on the TensorScatter op used for KV cache updates in attention mechanisms. See the tensor format notation defined in the TRT_Attention section above.

Attributes

update_form - STRING

"padded_bhnd" (default) or "packed_nhd". Sets the IO form of the update tensor.

mode - STRING

"linear" (default).

axis - INT

Scatter axis in the past_cache tensor. Must be -2 (default), which is the sequence-length dimension in the [B, H, N, D] layout.

Inputs (2-4)

past_cache (index 0)

4D tensor of shape [B, H, N, D].

update (index 1)

4D [B, H, N, D] for padded form or 3D [T, H, D] for packed form.

write_indices (optional, index 2) - INT32

1D tensor of shape [batchSize].

update_lengths (optional, index 3) - INT32

1D tensor of shape [batchSize + 1] containing cumulative token counts for the update sequence.

Outputs

Y

Restrictions

• past_cache must always be 4D [B, H, N, D].
• update must be 4D [B, H, N, D] for padded form or 3D [T, H, D] for packed form.