| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Rotary Positional Embeddings.""" |
| 4 | |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from .bailing_mrope import BailingMRotaryEmbedding |
| 10 | from .base import RotaryEmbedding |
| 11 | from .deepseek_scaling_rope import ( |
| 12 | DeepseekScalingRotaryEmbedding, |
| 13 | DeepseekV4ScalingRotaryEmbedding, |
| 14 | ) |
| 15 | from .dual_chunk_rope import DualChunkRotaryEmbedding |
| 16 | from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding |
| 17 | from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding |
| 18 | from .fope import FourierRotaryEmbedding |
| 19 | from .gemma4_rope import Gemma4RotaryEmbedding |
| 20 | from .linear_scaling_rope import LinearScalingRotaryEmbedding |
| 21 | from .llama3_rope import Llama3RotaryEmbedding |
| 22 | from .llama4_vision_rope import Llama4VisionRotaryEmbedding |
| 23 | from .mrope import MRotaryEmbedding |
| 24 | from .mrope_interleaved import MRotaryEmbeddingInterleaved |
| 25 | from .ntk_scaling_rope import NTKScalingRotaryEmbedding |
| 26 | from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding |
| 27 | from .yarn_scaling_rope import YaRNScalingRotaryEmbedding |
| 28 | |
| 29 | _ROPE_DICT: dict[tuple[Any, ...], RotaryEmbedding] = {} |
| 30 | |
| 31 | |
| 32 | def get_rope( |
| 33 | head_size: int, |
| 34 | max_position: int, |
| 35 | is_neox_style: bool = True, |
| 36 | rope_parameters: dict[str, Any] | None = None, |
| 37 | dtype: torch.dtype | None = None, |
| 38 | dual_chunk_attention_config: dict[str, Any] | None = None, |
| 39 | ) -> RotaryEmbedding: |
| 40 | if dtype is None: |
| 41 | dtype = torch.get_default_dtype() |
| 42 | if rope_parameters is not None: |
| 43 | # Transforms every value that is a list into a tuple for caching calls |
| 44 | rope_parameters_tuple = { |
| 45 | k: tuple(v) if isinstance(v, list) else v |
| 46 | for k, v in rope_parameters.items() |
| 47 | } |
| 48 | rope_parameters_args = tuple(rope_parameters_tuple.items()) |
| 49 | else: |
| 50 | rope_parameters_args = None |
| 51 | |
| 52 | if dual_chunk_attention_config is not None: |
| 53 | dual_chunk_attention_tuple = { |
| 54 | k: tuple(v) if isinstance(v, list) else v |
| 55 | for k, v in dual_chunk_attention_config.items() |
| 56 | if k != "sparse_attention_config" |
| 57 | } |
| 58 | dual_chunk_attention_args = tuple(dual_chunk_attention_tuple.items()) |
| 59 | else: |
| 60 | dual_chunk_attention_args = None |
| 61 | |
| 62 | rope_parameters = rope_parameters or {} |
| 63 | base = rope_parameters.get("rope_theta", 10000) |
| 64 | scaling_type = rope_parameters.get("rope_type", "default") |
| 65 | if rotary_dim := rope_parameters.get("rope_dim", None): |
| 66 | pass |
| 67 | else: |
| 68 | partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0) |
| 69 | if partial_rotary_factor <= 0.0 or partial_rotary_factor > 1.0: |
| 70 | raise ValueError(f"{partial_rotary_factor=} must be between 0.0 and 1.0") |
| 71 | rotary_dim = int(head_size * partial_rotary_factor) |
| 72 | |
| 73 | key = ( |
| 74 | head_size, |
| 75 | rotary_dim, |
| 76 | max_position, |
| 77 | is_neox_style, |
| 78 | rope_parameters_args, |
| 79 | dual_chunk_attention_args, |
| 80 | dtype, |
| 81 | ) |
| 82 | if key in _ROPE_DICT: |
| 83 | return _ROPE_DICT[key] |
| 84 | |
| 85 | if dual_chunk_attention_config is not None: |
| 86 | extra_kwargs = { |
| 87 | k: v |
| 88 | for k, v in dual_chunk_attention_config.items() |
| 89 | if k in ("chunk_size", "local_size") |
| 90 | } |
| 91 | rotary_emb = DualChunkRotaryEmbedding( |
| 92 | head_size, |
| 93 | rotary_dim, |
| 94 | max_position, |
| 95 | base, |
| 96 | is_neox_style, |
| 97 | dtype, |
| 98 | **extra_kwargs, |
| 99 | ) |
| 100 | elif scaling_type == "bailing_mrope": |
| 101 | rotary_emb = BailingMRotaryEmbedding( |
| 102 | head_size, |
| 103 | rotary_dim, |
| 104 | max_position, |
| 105 | base, |
| 106 | is_neox_style, |
| 107 | dtype, |
| 108 | mrope_section=rope_parameters["mrope_section"], |
| 109 | ) |
| 110 | elif scaling_type == "default": |
| 111 | if "mrope_section" in rope_parameters: |
| 112 | rotary_emb = MRotaryEmbedding( |
| 113 | head_size, |
| 114 | rotary_dim, |
| 115 | max_position, |
| 116 | base, |
| 117 | is_neox_style, |
| 118 | dtype, |
| 119 | mrope_section=rope_parameters["mrope_section"], |
| 120 | mrope_interleaved=rope_parameters.get("mrope_interleaved", False), |
| 121 | ) |
| 122 | elif "use_fope" in rope_parameters and rope_parameters["use_fope"]: |
| 123 | extra_kwargs = { |
| 124 | k: v |
| 125 | for k, v in rope_parameters.items() |
| 126 | if k |
| 127 | in ( |
| 128 | "num_key_value_heads", |
| 129 | "num_inv_freq", |
| 130 | "fope_sep_head", |
| 131 | "fope_init_factor", |
| 132 | ) |
| 133 | } |
| 134 | extra_kwargs["init_cache"] = False |
| 135 | rotary_emb = FourierRotaryEmbedding( |
| 136 | head_size, |
| 137 | rotary_dim, |
| 138 | max_position, |
| 139 | base, |
| 140 | is_neox_style, |
| 141 | dtype, |
| 142 | **extra_kwargs, |
| 143 | ) |
| 144 | else: |
| 145 | rotary_emb = RotaryEmbedding( |
| 146 | head_size, |
| 147 | rotary_dim, |
| 148 | max_position, |
| 149 | base, |
| 150 | is_neox_style, |
| 151 | dtype, |
| 152 | ) |
| 153 | elif scaling_type == "proportional": |
| 154 | # Proportional RoPE is used by Gemma4 for global (full) attention. |
| 155 | # Gemma4 uses a sparse/fractional RoPE with cross-mixing between halves. |
| 156 | rotary_emb = Gemma4RotaryEmbedding( |
| 157 | head_size, |
| 158 | rotary_dim, |
| 159 | max_position, |
| 160 | base, |
| 161 | is_neox_style, |
| 162 | dtype, |
| 163 | ) |
| 164 | elif scaling_type == "llama3": |
| 165 | scaling_factor = rope_parameters["factor"] |
| 166 | low_freq_factor = rope_parameters["low_freq_factor"] |
| 167 | high_freq_factor = rope_parameters["high_freq_factor"] |
| 168 | original_max_position = rope_parameters["original_max_position_embeddings"] |
| 169 | rotary_emb = Llama3RotaryEmbedding( |
| 170 | head_size, |
| 171 | rotary_dim, |
| 172 | max_position, |
| 173 | base, |
| 174 | is_neox_style, |
| 175 | dtype, |
| 176 | scaling_factor, |
| 177 | low_freq_factor, |
| 178 | high_freq_factor, |
| 179 | original_max_position, |
| 180 | ) |
| 181 | elif scaling_type == "mllama4": |
| 182 | rotary_emb = Llama4VisionRotaryEmbedding( |
| 183 | head_size, rotary_dim, max_position, base, is_neox_style, dtype |
| 184 | ) |
| 185 | elif scaling_type == "linear": |
| 186 | scaling_factor = rope_parameters["factor"] |
| 187 | rotary_emb = LinearScalingRotaryEmbedding( |
| 188 | head_size, |
| 189 | rotary_dim, |
| 190 | max_position, |
| 191 | base, |
| 192 | is_neox_style, |
| 193 | scaling_factor, |
| 194 | dtype, |
| 195 | ) |
| 196 | elif scaling_type == "ntk": |
| 197 | scaling_factor = rope_parameters["factor"] |
| 198 | mixed_b = rope_parameters.get("mixed_b") |
| 199 | rotary_emb = NTKScalingRotaryEmbedding( |
| 200 | head_size, |
| 201 | rotary_dim, |
| 202 | max_position, |
| 203 | base, |
| 204 | is_neox_style, |
| 205 | scaling_factor, |
| 206 | dtype, |
| 207 | mixed_b, |
| 208 | ) |
| 209 | elif scaling_type == "dynamic": |
| 210 | if "alpha" in rope_parameters: |
| 211 | scaling_alpha = rope_parameters["alpha"] |
| 212 | rotary_emb = DynamicNTKAlphaRotaryEmbedding( |
| 213 | head_size, |
| 214 | rotary_dim, |
| 215 | max_position, |
| 216 | base, |
| 217 | is_neox_style, |
| 218 | scaling_alpha, |
| 219 | dtype, |
| 220 | ) |
| 221 | elif "factor" in rope_parameters: |
| 222 | scaling_factor = rope_parameters["factor"] |
| 223 | max_trained_positions = rope_parameters.get( |
| 224 | "max_trained_positions", max_position |
| 225 | ) |
| 226 | rotary_emb = DynamicNTKScalingRotaryEmbedding( |
| 227 | head_size, |
| 228 | rotary_dim, |
| 229 | max_position, |
| 230 | max_trained_positions, |
| 231 | base, |
| 232 | is_neox_style, |
| 233 | scaling_factor, |
| 234 | dtype, |
| 235 | ) |
| 236 | else: |
| 237 | raise ValueError( |
| 238 | "Dynamic rope scaling must contain either 'alpha' or 'factor' field" |
| 239 | ) |
| 240 | elif scaling_type == "yarn": |
| 241 | scaling_factor = rope_parameters["factor"] |
| 242 | original_max_position = rope_parameters["original_max_position_embeddings"] |
| 243 | extra_kwargs = { |
| 244 | k: v |
| 245 | for k, v in rope_parameters.items() |
| 246 | if k |
| 247 | in ( |
| 248 | "beta_fast", |
| 249 | "beta_slow", |
| 250 | "mscale", |
| 251 | "mscale_all_dim", |
| 252 | "attention_factor", |
| 253 | "truncate", |
| 254 | ) |
| 255 | } |
| 256 | if "mrope_section" in rope_parameters: |
| 257 | rotary_emb = MRotaryEmbedding( |
| 258 | head_size, |
| 259 | rotary_dim, |
| 260 | original_max_position, |
| 261 | base, |
| 262 | is_neox_style, |
| 263 | dtype, |
| 264 | mrope_section=rope_parameters["mrope_section"], |
| 265 | mrope_interleaved=rope_parameters.get("mrope_interleaved", False), |
| 266 | scaling_factor=scaling_factor, |
| 267 | **extra_kwargs, |
| 268 | ) |
| 269 | else: |
| 270 | rotary_emb = YaRNScalingRotaryEmbedding( |
| 271 | head_size, |
| 272 | rotary_dim, |
| 273 | original_max_position, |
| 274 | base, |
| 275 | is_neox_style, |
| 276 | scaling_factor, |
| 277 | dtype, |
| 278 | **extra_kwargs, |
| 279 | ) |
| 280 | elif scaling_type in ["deepseek_yarn", "deepseek_llama_scaling"]: |
| 281 | scaling_factor = rope_parameters["factor"] |
| 282 | original_max_position = rope_parameters["original_max_position_embeddings"] |
| 283 | # assert max_position == original_max_position * scaling_factor |
| 284 | extra_kwargs = { |
| 285 | k: v |
| 286 | for k, v in rope_parameters.items() |
| 287 | if k |
| 288 | in ( |
| 289 | "beta_fast", |
| 290 | "beta_slow", |
| 291 | "mscale", |
| 292 | "mscale_all_dim", |
| 293 | ) |
| 294 | } |
| 295 | if rope_parameters.get("is_deepseek_v4", False): |
| 296 | cls = DeepseekV4ScalingRotaryEmbedding |
| 297 | else: |
| 298 | cls = DeepseekScalingRotaryEmbedding |
| 299 | rotary_emb = cls( |
| 300 | head_size, |
| 301 | rotary_dim, |
| 302 | original_max_position, |
| 303 | base, |
| 304 | is_neox_style, |
| 305 | scaling_factor, |
| 306 | dtype, |
| 307 | **extra_kwargs, |
| 308 | ) |
| 309 | elif scaling_type == "longrope": |
| 310 | short_factor = rope_parameters["short_factor"] |
| 311 | long_factor = rope_parameters["long_factor"] |
| 312 | original_max_position = rope_parameters["original_max_position_embeddings"] |
| 313 | extra_kwargs = { |
| 314 | k: v |
| 315 | for k, v in rope_parameters.items() |
| 316 | if k in ("short_mscale", "long_mscale") |
| 317 | } |
| 318 | rotary_emb = Phi3LongRoPEScaledRotaryEmbedding( |
| 319 | head_size, |
| 320 | rotary_dim, |
| 321 | max_position, |
| 322 | original_max_position, |
| 323 | base, |
| 324 | is_neox_style, |
| 325 | dtype, |
| 326 | short_factor, |
| 327 | long_factor, |
| 328 | **extra_kwargs, |
| 329 | ) |
| 330 | elif scaling_type == "openpangu": |
| 331 | mrope_interleaved = rope_parameters.get("mrope_interleaved", False) |
| 332 | if "mrope_section" in rope_parameters and mrope_interleaved: |
| 333 | rotary_emb = MRotaryEmbeddingInterleaved( |
| 334 | head_size, |
| 335 | rotary_dim, |
| 336 | max_position, |
| 337 | base, |
| 338 | is_neox_style, |
| 339 | dtype, |
| 340 | mrope_section=rope_parameters["mrope_section"], |
| 341 | mrope_interleaved=mrope_interleaved, |
| 342 | ) |
| 343 | else: |
| 344 | raise ValueError("Pangu mrope lacks necessary parameters.") |
| 345 | else: |
| 346 | raise ValueError(f"Unknown RoPE scaling type {scaling_type}") |
| 347 | _ROPE_DICT[key] = rotary_emb |
| 348 | return rotary_emb |