| 30 | |
| 31 | |
| 32 | def build_alibi_bias( |
| 33 | n_heads, seq_len, full=False, alibi_bias_max=8, dtype=torch.float32 |
| 34 | ): |
| 35 | alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32).view(1, 1, 1, seq_len) |
| 36 | if full: |
| 37 | alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32).view( |
| 38 | 1, 1, seq_len, 1 |
| 39 | ) |
| 40 | alibi_bias = alibi_bias.abs().mul(-1) |
| 41 | slopes = gen_slopes(n_heads, alibi_bias_max) |
| 42 | alibi_bias = alibi_bias * slopes |
| 43 | slopes = slopes.squeeze(0).squeeze(-1).squeeze(-1) |
| 44 | return slopes.to(dtype=dtype), alibi_bias.to(dtype=dtype) |
| 45 | |
| 46 | |
| 47 | def _cast_if_autocast_enabled(tensor): |