| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from .base import RotaryEmbedding |
| 9 | |
| 10 | |
| 11 | class Llama3RotaryEmbedding(RotaryEmbedding): |
| 12 | def __init__( |
| 13 | self, |
| 14 | head_size: int, |
| 15 | rotary_dim: int, |
| 16 | max_position_embeddings: int, |
| 17 | base: float, |
| 18 | is_neox_style: bool, |
| 19 | dtype: torch.dtype, |
| 20 | scaling_factor: float, |
| 21 | low_freq_factor: float, |
| 22 | high_freq_factor: float, |
| 23 | orig_max_position: int, |
| 24 | ) -> None: |
| 25 | self.scaling_factor = scaling_factor |
| 26 | self.low_freq_factor = low_freq_factor |
| 27 | self.high_freq_factor = high_freq_factor |
| 28 | self.orig_max_position = orig_max_position |
| 29 | super().__init__( |
| 30 | head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype |
| 31 | ) |
| 32 | |
| 33 | def _compute_inv_freq(self, base: float) -> torch.Tensor: |
| 34 | inv_freqs = super()._compute_inv_freq(base) |
| 35 | low_freq_wavelen = self.orig_max_position / self.low_freq_factor |
| 36 | high_freq_wavelen = self.orig_max_position / self.high_freq_factor |
| 37 | |
| 38 | wave_len = 2 * math.pi / inv_freqs |
| 39 | if self.low_freq_factor != self.high_freq_factor: |
| 40 | smooth = (self.orig_max_position / wave_len - self.low_freq_factor) / ( |
| 41 | self.high_freq_factor - self.low_freq_factor |
| 42 | ) |
| 43 | else: |
| 44 | smooth = 0 |
| 45 | new_freqs = torch.where( |
| 46 | wave_len < high_freq_wavelen, |
| 47 | inv_freqs, |
| 48 | torch.where( |
| 49 | wave_len > low_freq_wavelen, |
| 50 | inv_freqs / self.scaling_factor, |
| 51 | (1 - smooth) * inv_freqs / self.scaling_factor + smooth * inv_freqs, |
| 52 | ), |
| 53 | ) |
| 54 | return new_freqs |