Forum
Login
Register

Welcome to our Support Center
UpdatedAugust 22, 2025

PackedMultiHeadAttention

Description

This is the packed version of MultiHeadAttention.

 

 

Sequences in one batch usually don’t have same length and they are padded to have same length, e.g., below is a batch with 3 sequences and * is padding token. Sequence_0: 0, 1*, 2*, 3* Sequence_1: 4, 5, 6*, 7* Sequence_2: 8, 9, 10, 11

PackedMultiHeadAttention is designed to takes in packed input, i.e., only the real tokens without padding. An input as above will be packed into 3 tensors like below:

    • query ([q0, q4, q5, q8, q9, q10, q11])
    • key ([k0, k4, k5, k8, k9, k10, k11])
    • value ([v0, v4, v5, v8, v9, v10, v11])
    • token_offset: 0, 4, 5, 8, 9, 10, 11, 1*, 2*, 3*, 6*, 7*
    • cumulative_sequence_length: 0, 1, 1+2, 1+2+4

The query, key and value tensors contain result of hidden embedding of real tokens after input projections. Token_offset records the offset of token in the unpacked input. cumulative_sequence_length records cumulated length of each sequence length.

The operator only supports BERT like model with padding on right now.

Input parameters

 

specified_outputs_name :Β array, this parameter lets you manually assign custom names to the output tensors of a node.

Β Graphs in :Β cluster, ONNX model architecture.

query (heterogeneous)Β – T : object, query with shape (token_count, hidden_size) or packed qkv with shape (token_count, num_heads, 3, head_size).
key (optional, heterogeneous) – T : object, key with shape (token_count, hidden_size).
value (optional, heterogeneous) – T : object, value with shape (token_count, v_hidden_size).
bias (optional, heterogeneous) – T : object, bias tensor with shape (hidden_size + hidden_size + v_hidden_size) from input projection.
token_offset (heterogeneous) – M : object, offset of each token before packing, with shape (batch_size, sequence_length).
cumulative_sequence_length (heterogeneous) – M : object, aΒ tensor with shape (batch_size + 1). It specifies the cumulative sequence length.
relative_position_bias (optional, heterogeneous)Β – T : object, it specifies the additional bias to QxK’. The shape is (batch_size or 1, num_heads or 1, sequence_length, sequence_length).

Β Parameters :Β cluster,

mask_filter_value : float, the value to be filled in the attention mask.
Default value β€œ-10000”.
num_headsΒ :Β integer, number of attention heads.
Default value β€œ0”.
scale : float, custom scale will be used if specified.
Default value β€œ0”.
Β training?Β :Β boolean, whether the layer is in training mode (can store data for backward).
Default value β€œTrue”.
Β lda coeff :Β float, defines the coefficient by which the loss derivative will be multiplied before being sent to the previous layer (since during the backward run we go backwards).
Default value β€œ1”.

Β name (optional) :Β string,Β name of the node.

Output parameters

output (heterogeneous)Β – T : object, output tensor with shape (token_count, v_hidden_size).

Type Constraints

TΒ in (tensor(float),Β tensor(float16)) : Constrain input and output to float tensors.

MΒ in (tensor(int32)) : Constrain mask, offset and sequence length to integer types.

Example

All these exemples are snippets PNG, you can drop these Snippet onto the block diagram and get the depicted code added to your VI (Do not forget to install Deep Learning library to run it).
Table of Contents