Function bert_attention

tensorrt_llm/functional.py:4487–4676

Add an operation that performs the multi-head attention in BERT. The multi-head attention (MHA) is the sequence of a batched matmul, a softmax and a batched matmul as described in https://arxiv.org/abs/1706.03762. That function adds an operation that performs those computations

(tensor: Tensor,
                   input_lengths: Tensor,
                   num_heads: int,
                   head_size: int,
                   q_scaling: float,
                   relative_attention: bool = False,
                   relative_attention_bias: Tensor = None,
                   max_distance: int = 0,
                   max_input_length: Tensor = None,
                   sage_attn: bool = False,
                   sage_attn_q_block_size: int = 0,
                   sage_attn_k_block_size: int = 0,
                   sage_attn_v_block_size: int = 0,
                   cp_group: list[int] = None,
                   cp_size: int = 1,
                   cp_rank: int = 0)

def bert_attention(tensor: Tensor,
                   input_lengths: Tensor,
                   num_heads: int,
                   head_size: int,
                   q_scaling: float,
                   relative_attention: bool = False,
                   relative_attention_bias: Tensor = None,
                   max_distance: int = 0,
                   max_input_length: Tensor = None,
                   sage_attn: bool = False,
                   sage_attn_q_block_size: int = 0,
                   sage_attn_k_block_size: int = 0,
                   sage_attn_v_block_size: int = 0,
                   cp_group: list[int] = None,
                   cp_size: int = 1,
                   cp_rank: int = 0) -> Tuple[Tensor]:
    '''
    Add an operation that performs the multi-head attention in BERT.

    The multi-head attention (MHA) is the sequence of a batched matmul, a
    softmax and a batched matmul as described in
    https://arxiv.org/abs/1706.03762. That function adds an operation that
    performs those computations using a single GPU kernel.

    The input tensor contains the Q, K and V elements. It is a 2D tensor and
    its shape is '[sum_of_tokens, 3*hidden_dim]' where the 'sum_of_tokens' is
    the sum of the sequence lengths in the batch.

    In MHA, the output of the Q*K^T product is scaled by a constant value that
    is computed as:

        1.f / (q_scaling * sqrt(head_size)).

    That 'q_scaling' constant is the last argument of that function.

    That layer is implemented using a plugin (see bertAttentionPlugin).

    Parameters:
        tensor : Tensor
            The QKV input tensor.

        input_lengths : Tensor
            The length of each sequence. It is a 1D tensor of size 'batch_size'.

        num_heads : int
            The number of heads.

        head_size : int
            The size of each head.

        q_scaling : float
            The factor to compute the scaling factor to scale the output of the
            'Q*K^T' product.

        relative_attention: bool = False
            If enable relative attention.

        relative_attention_bias: Tensor = None