utils

Eagle model utils.

Functions

expand_mask

Expands attention_mask from [bsz, seq_len] to [bsz, 1, tgt_seq_len, src_seq_len].

make_causal_mask

Make causal mask used for bi-directional self-attention.
expand_mask(mask, dtype, tgt_len=None)

Expands attention_mask from [bsz, seq_len] to [bsz, 1, tgt_seq_len, src_seq_len].

Parameters:

  • mask (Tensor)

  • dtype (dtype)

  • tgt_len (int | None)

make_causal_mask(input_ids_shape, dtype, device, past_key_values_length=0)

Make causal mask used for bi-directional self-attention.

Parameters:

  • input_ids_shape (Size)

  • dtype (dtype)

  • device (device)

  • past_key_values_length (int)