Function expand_mask

tensorrt_llm/functional.py:6467–6496

Expand an attention mask. That function adds the sequence of operations to expand from a tensor of shape '[batch_size, src_seq_len]' to a tensor of shape '[batch_size, 1, tgt_seq_len, src_seq_len]'. It can be used to create the mask applied to the Q*K^T product before the softm

(mask: Tensor, tgt_len: Optional[Tensor] = None)

def expand_mask(mask: Tensor, tgt_len: Optional[Tensor] = None) -> Tensor:
    '''
    Expand an attention mask.

    That function adds the sequence of operations to expand from a tensor of
    shape '[batch_size, src_seq_len]' to a tensor of shape
    '[batch_size, 1, tgt_seq_len, src_seq_len]'. It can be used to create the
    mask applied to the Q*K^T product before the softmax operation in the
    multi-head attention block.

    Parameters:
        mask : Tensor
            The input mask

        tgt_len : Optional[Tensor]
            The dimension of the 3rd dimension in the output tensor. If None,
            the 2nd dimension of the input is used.

    Returns:
        The tensor created by that sequence of operations.
    '''
    bsz = shape(mask, 0)
    src_len = shape(mask, 1)
    tgt_len = tgt_len if tgt_len is not None else src_len

    mask = mask.view(concat([bsz, 1, 1, src_len]))

    mask = expand(mask, concat([bsz, 1, tgt_len, src_len]))
    mask = where(mask == 0, float('-inf'), 0.0)
    return mask


def gather_last_token_logits(hidden_states: Tensor, last_token_ids: Tensor,