| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepseekV4 rotary embedding initialization.""" |
| 4 | |
| 5 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 6 | from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding |
| 7 | |
| 8 | |
| 9 | def build_deepseek_v4_rope( |
| 10 | config, |
| 11 | *, |
| 12 | head_dim: int, |
| 13 | rope_head_dim: int, |
| 14 | max_position_embeddings: int, |
| 15 | compress_ratio: int, |
| 16 | ) -> RotaryEmbedding: |
| 17 | rope_parameters = config.rope_parameters |
| 18 | # Newer checkpoints nest per-layer-type rope dicts ({"main", "compress"}); |
| 19 | # older ones ship a single flat dict shared by all layer types. |
| 20 | if isinstance(rope_parameters.get("main"), dict) and isinstance( |
| 21 | rope_parameters.get("compress"), dict |
| 22 | ): |
| 23 | key = "compress" if compress_ratio > 0 else "main" |
| 24 | rope_parameters = dict(rope_parameters[key]) |
| 25 | else: |
| 26 | rope_parameters = dict(rope_parameters) |
| 27 | |
| 28 | rope_parameters["rope_theta"] = ( |
| 29 | config.compress_rope_theta if compress_ratio > 0 else config.rope_theta |
| 30 | ) |
| 31 | if compress_ratio > 0 and rope_parameters["rope_type"] != "default": |
| 32 | # YaRN applies to all compressed-KV layers (compress_ratio > 0); only |
| 33 | # pure sliding-window layers (compress_ratio == 0) use plain RoPE. |
| 34 | rope_parameters["rope_type"] = ( |
| 35 | "deepseek_llama_scaling" |
| 36 | if rope_parameters.get("attention_factor") == 1.0 |
| 37 | else "deepseek_yarn" |
| 38 | ) |
| 39 | else: |
| 40 | # Sliding-window layers use plain RoPE (theta=rope_theta, no YaRN). |
| 41 | rope_parameters["rope_type"] = "deepseek_yarn" |
| 42 | rope_parameters["factor"] = 1.0 |
| 43 | rope_parameters["original_max_position_embeddings"] = max_position_embeddings |
| 44 | rope_parameters["mscale"] = 0 # Disable mscale |
| 45 | rope_parameters["mscale_all_dim"] = 0 # Disable mscale |
| 46 | rope_parameters["is_deepseek_v4"] = True |
| 47 | rope_parameters["rope_dim"] = rope_head_dim |
| 48 | return get_rope( |
| 49 | head_dim, |
| 50 | max_position=max_position_embeddings, |
| 51 | rope_parameters=rope_parameters, |
| 52 | is_neox_style=False, |
| 53 | ) |