| 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 vllm.platforms import current_platform |
| 9 | from vllm.utils.flashinfer import has_flashinfer |
| 10 | |
| 11 | from .base import RotaryEmbeddingBase |
| 12 | from .common import ( |
| 13 | rotate_gptj, |
| 14 | rotate_neox, |
| 15 | yarn_find_correction_range, |
| 16 | yarn_linear_ramp_mask, |
| 17 | ) |
| 18 | |
| 19 | |
| 20 | def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: |
| 21 | if scale <= 1: |
| 22 | return 1.0 |
| 23 | return 0.1 * mscale * math.log(scale) + 1.0 |
| 24 | |
| 25 | |
| 26 | class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase): |
| 27 | """RotaryEmbedding extended with YaRN method. |
| 28 | |
| 29 | Credits to Peng et al. github.com/jquesnelle/yarn |
| 30 | """ |
| 31 | |
| 32 | def __init__( |
| 33 | self, |
| 34 | head_size: int, |
| 35 | rotary_dim: int, |
| 36 | max_position_embeddings: int, |
| 37 | base: float, |
| 38 | is_neox_style: bool, |
| 39 | scaling_factor: float, |
| 40 | dtype: torch.dtype, |
| 41 | *, |
| 42 | beta_fast: int = 32, |
| 43 | beta_slow: int = 1, |
| 44 | mscale: float = 1, |
| 45 | mscale_all_dim: float = 0, |
| 46 | init_cache: bool = True, |
| 47 | ) -> None: |
| 48 | self.scaling_factor = scaling_factor |
| 49 | self.beta_fast = beta_fast |
| 50 | self.beta_slow = beta_slow |
| 51 | # Get n-d magnitude scaling corrected for interpolation. |
| 52 | self.mscale = float( |
| 53 | yarn_get_mscale(self.scaling_factor, float(mscale)) |
| 54 | / yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) |
| 55 | ) |
| 56 | self.use_flashinfer = ( |
| 57 | self.enabled() |
| 58 | and dtype in (torch.float16, torch.bfloat16) |
| 59 | and current_platform.is_cuda() |
| 60 | and has_flashinfer() |
| 61 | and head_size in [64, 128, 256, 512] |
| 62 | ) |
| 63 | super().__init__( |
| 64 | head_size, |
| 65 | rotary_dim, |
| 66 | max_position_embeddings, |
| 67 | base, |
| 68 | is_neox_style, |
| 69 | dtype, |
| 70 | init_cache=init_cache, |
| 71 | ) |
| 72 | |
| 73 | def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor: |
| 74 | pos_freqs = self.base ** ( |
| 75 | torch.arange( |
| 76 | 0, |
| 77 | self.rotary_dim, |
| 78 | 2, |
| 79 | dtype=torch.float, |
| 80 | ) |
| 81 | / self.rotary_dim |
| 82 | ) |
| 83 | inv_freq_extrapolation = 1.0 / pos_freqs |
| 84 | inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs) |
| 85 | |
| 86 | low, high = yarn_find_correction_range( |
| 87 | self.beta_fast, |
| 88 | self.beta_slow, |
| 89 | self.rotary_dim, |
| 90 | self.base, |
| 91 | self.max_position_embeddings, |
| 92 | ) |
| 93 | # Get n-d rotational scaling corrected for extrapolation |
| 94 | inv_freq_mask = 1 - yarn_linear_ramp_mask( |
| 95 | low, high, self.rotary_dim // 2, dtype=torch.float |
| 96 | ) |
| 97 | inv_freq = ( |
| 98 | inv_freq_interpolation * (1 - inv_freq_mask) |
| 99 | + inv_freq_extrapolation * inv_freq_mask |
| 100 | ) |
| 101 | return inv_freq |
| 102 | |
| 103 | def _compute_cos_sin_cache(self) -> torch.Tensor: |
| 104 | inv_freq = self._compute_inv_freq(self.scaling_factor) |
| 105 | t = torch.arange( |
| 106 | self.max_position_embeddings * self.scaling_factor, |
| 107 | dtype=torch.float32, |
| 108 | ) |
| 109 | freqs = torch.einsum("i,j -> ij", t, inv_freq) |
| 110 | cos = freqs.cos() * self.mscale |
| 111 | sin = freqs.sin() * self.mscale |
| 112 | cache = torch.cat((cos, sin), dim=-1) |
| 113 | return cache |
| 114 | |
| 115 | def forward_native( |
| 116 | self, |
| 117 | positions: torch.Tensor, |
| 118 | query: torch.Tensor, |
| 119 | key: torch.Tensor | None = None, |
| 120 | offsets: torch.Tensor | None = None, |
| 121 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 122 | """PyTorch-native implementation equivalent to forward().""" |
| 123 | assert key is not None |
| 124 | return self.forward_static( |
| 125 | positions, |
| 126 | query, |
| 127 | key, |
| 128 | self.head_size, |
| 129 | self.rotary_dim, |
| 130 | self.cos_sin_cache, |
| 131 | self.is_neox_style, |
| 132 | offsets, |
| 133 | ) |
| 134 | |
| 135 | @staticmethod |
| 136 | def forward_static( |
| 137 | positions: torch.Tensor, |
| 138 | query: torch.Tensor, |
| 139 | key: torch.Tensor | None, |
| 140 | head_size: int, |
| 141 | rotary_dim: int, |
| 142 | cos_sin_cache: torch.Tensor, |
| 143 | is_neox_style: bool, |
| 144 | offsets: torch.Tensor | None = None, |
| 145 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 146 | """A static implementation of forward().""" |
| 147 | assert key is not None |
| 148 | query_rot = query[..., :rotary_dim] |
| 149 | key_rot = key[..., :rotary_dim] |
| 150 | if rotary_dim < head_size: |
| 151 | query_pass = query[..., rotary_dim:] |
| 152 | key_pass = key[..., rotary_dim:] |
| 153 | |
| 154 | cos_sin = cos_sin_cache[ |
| 155 | torch.add(positions, offsets) if offsets is not None else positions |
| 156 | ] |
| 157 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 158 | if is_neox_style: |
| 159 | cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) |
| 160 | sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2) |
| 161 | else: |
| 162 | cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) |
| 163 | sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) |
| 164 | |
| 165 | rotate_fn = rotate_neox if is_neox_style else rotate_gptj |
| 166 | query_rot = query_rot * cos + rotate_fn(query_rot) * sin |
| 167 | key_rot = key_rot * cos + rotate_fn(key_rot) * sin |
| 168 | |
| 169 | if rotary_dim < head_size: |
| 170 | query = torch.cat((query_rot, query_pass), dim=-1) |
| 171 | key = torch.cat((key_rot, key_pass), dim=-1) |
| 172 | else: |
| 173 | query = query_rot |
| 174 | key = key_rot |
| 175 | return query, key |
| 176 | |
| 177 | def forward_xpu( |
| 178 | self, |
| 179 | positions: torch.Tensor, |
| 180 | query: torch.Tensor, |
| 181 | key: torch.Tensor | None = None, |
| 182 | offsets: torch.Tensor | None = None, |
| 183 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 184 | return torch.ops.vllm.xpu_ops_deepseek_scaling_rope( |
| 185 | positions, |
| 186 | query, |
| 187 | key, |
| 188 | offsets, |
| 189 | self._match_cos_sin_cache_dtype(query), |
| 190 | self.rotary_dim, |
| 191 | self.is_neox_style, |
| 192 | ) |
| 193 | |
| 194 | def forward_hip( |
| 195 | self, |
| 196 | positions: torch.Tensor, |
| 197 | query: torch.Tensor, |
| 198 | key: torch.Tensor | None = None, |
| 199 | offsets: torch.Tensor | None = None, |
| 200 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 201 | return self.forward_native(positions, query, key, offsets) |
| 202 | |
| 203 | def forward_cuda( |
| 204 | self, |
| 205 | positions: torch.Tensor, |
| 206 | query: torch.Tensor, |
| 207 | key: torch.Tensor | None = None, |
| 208 | offsets: torch.Tensor | None = None, |
| 209 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 210 | if self.use_flashinfer: |
| 211 | torch.ops.vllm.flashinfer_rotary_embedding( |
| 212 | torch.add(positions, offsets) if offsets is not None else positions, |
| 213 | query, |
| 214 | key, |
| 215 | self.head_size, |
| 216 | self.cos_sin_cache, |
| 217 | self.is_neox_style, |
| 218 | ) |
| 219 | return query, key |
| 220 | else: |
| 221 | return self.forward_native(positions, query, key, offsets) |
| 222 | |
| 223 | |
| 224 | class DeepseekV4ScalingRotaryEmbedding(DeepseekScalingRotaryEmbedding): |
| 225 | """RotaryEmbedding extended with YaRN method. |
| 226 | |
| 227 | Credits to Peng et al. github.com/jquesnelle/yarn |
| 228 | |
| 229 | Compared to DeepseekScalingRotaryEmbedding: |
| 230 | - Applies RoPE to the last rotary_dim |
| 231 | - The forward method requires an inverse parameter to indicate |
| 232 | whether to negate the sin |
| 233 | - Supports applying RoPE to query only (without key) |
| 234 | - cos_sin_cache stored as fp32 for higher precision RoPE |
| 235 | """ |
| 236 | |
| 237 | def __init__(self, *args, **kwargs): |
| 238 | # Avoid compute cache repeatedly |
| 239 | kwargs.pop("init_cache", None) |
| 240 | super().__init__(*args, **kwargs, init_cache=False) |
| 241 | cache_fp32 = self._compute_cos_sin_cache() |
| 242 | self.register_buffer("cos_sin_cache", cache_fp32, persistent=False) |
| 243 | |
| 244 | def _compute_cos_sin_cache(self) -> torch.Tensor: |
| 245 | inv_freq = self._compute_inv_freq(self.scaling_factor) |
| 246 | t = torch.arange( |
| 247 | self.max_position_embeddings * self.scaling_factor, |
| 248 | device=inv_freq.device, |
| 249 | dtype=torch.float32, |
| 250 | ) |
| 251 | freqs = torch.einsum("i,j -> ij", t, inv_freq) |
| 252 | cos = freqs.cos() * self.mscale |
| 253 | sin = freqs.sin() * self.mscale |
| 254 | cache = torch.cat((cos, sin), dim=-1) |
| 255 | return cache |
| 256 | |
| 257 | def forward_native( |
| 258 | self, |
| 259 | positions: torch.Tensor, |
| 260 | query: torch.Tensor, |
| 261 | key: torch.Tensor | None = None, |
| 262 | offsets: torch.Tensor | None = None, |
| 263 | inverse: bool = False, |
| 264 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 265 | """PyTorch-native implementation equivalent to forward().""" |
| 266 | head_size = query.size(-1) |
| 267 | query_rot = query[..., -self.rotary_dim :] |
| 268 | key_rot = key[..., -self.rotary_dim :] if key is not None else None |
| 269 | |
| 270 | if self.rotary_dim < head_size: |
| 271 | query_pass = query[..., : -self.rotary_dim] |
| 272 | key_pass = key[..., : -self.rotary_dim] if key is not None else None |
| 273 | |
| 274 | cos_sin = self.cos_sin_cache[ |
| 275 | torch.add(positions, offsets) if offsets is not None else positions |
| 276 | ] |
| 277 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 278 | if self.is_neox_style: |
| 279 | cos = torch.cat((cos, cos), dim=-1).unsqueeze(-2) |
| 280 | sin = torch.cat((sin, sin), dim=-1).unsqueeze(-2) |
| 281 | else: |
| 282 | cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) |
| 283 | sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) |
| 284 | if inverse: |
| 285 | sin = -sin |
| 286 | rotate_fn = rotate_neox if self.is_neox_style else rotate_gptj |
| 287 | orig_dtype = query.dtype |
| 288 | query_rot = (query_rot * cos + rotate_fn(query_rot) * sin).to(orig_dtype) |
| 289 | if key_rot is not None: |
| 290 | key_rot = (key_rot * cos + rotate_fn(key_rot) * sin).to(orig_dtype) |
| 291 | |
| 292 | if self.rotary_dim < head_size: |
| 293 | query = torch.cat((query_pass, query_rot), dim=-1) |
| 294 | key = torch.cat((key_pass, key_rot), dim=-1) if key is not None else None |
| 295 | else: |
| 296 | query = query_rot |
| 297 | key = key_rot |
| 298 | |
| 299 | return query, key |
| 300 | |
| 301 | def forward_hip( |
| 302 | self, |
| 303 | positions: torch.Tensor, |
| 304 | query: torch.Tensor, |
| 305 | key: torch.Tensor | None = None, |
| 306 | offsets: torch.Tensor | None = None, |
| 307 | inverse: bool = False, |
| 308 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 309 | return self.forward_native(positions, query, key, offsets) |
| 310 | |
| 311 | def forward_cuda( |
| 312 | self, |
| 313 | positions: torch.Tensor, |
| 314 | query: torch.Tensor, |
| 315 | key: torch.Tensor | None = None, |
| 316 | offsets: torch.Tensor | None = None, |
| 317 | inverse: bool = False, |
| 318 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 319 | from vllm import _custom_ops as ops |
| 320 | |
| 321 | # The indexer and attention have different head_dim, |
| 322 | # we obtain the corresponding head_dim via the query. |
| 323 | head_size = query.size(-1) |
| 324 | rope_dim_offset = head_size - self.rotary_dim |
| 325 | # ops.rotary_embedding() is an in-place operation |
| 326 | # that updates the query and key tensors. |
| 327 | ops.rotary_embedding( |
| 328 | torch.add(positions, offsets) if offsets is not None else positions, |
| 329 | query, |
| 330 | key, |
| 331 | head_size, |
| 332 | self.cos_sin_cache, |
| 333 | self.is_neox_style, |
| 334 | rope_dim_offset=rope_dim_offset, |
| 335 | inverse=inverse, |
| 336 | ) |
| 337 | return query, key |