A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 07:15 UTC
[DSV4] Refactor RoPE initialization (#44262)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 1, 2026 at 18:26 UTC
517e74a9644fd8a21805fa61855ca787b3063308
3 files changed
+50
-40
vllm/models/deepseek_v4/amd/model.py
+7
-20
@@ -30,7 +30,6 @@ from vllm.model_executor.layers.mhc import (
30
MHCPreOp,
31
)
32
from vllm.model_executor.layers.quantization import QuantizationConfig
33
-from vllm.model_executor.layers.rotary_embedding import get_rope
33
from vllm.model_executor.layers.vocab_parallel_embedding import (
34
ParallelLMHead,
35
VocabParallelEmbedding,
@@ -50,6 +49,7 @@ from vllm.models.deepseek_v4.attention import (
49
DeepseekV4Indexer,
50
DeepseekV4MLA,
51
)
52
+from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
53
from vllm.platforms import current_platform
54
from vllm.sequence import IntermediateTensors
55
from vllm.utils.import_utils import has_tilelang
@@ -314,25 +314,12 @@ class DeepseekV4Attention(nn.Module):
314
self.rope_parameters = config.rope_scaling
315
316
# Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it)
317
- rope_parameters = config.rope_parameters
318
- rope_parameters["rope_theta"] = (
319
- config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta
320
- )
321
- if config.rope_parameters["rope_type"] != "default":
322
- config.rope_parameters["rope_type"] = (
323
- "deepseek_yarn"
324
- if config.rope_parameters.get("apply_yarn_scaling", True)
325
- else "deepseek_llama_scaling"
326
- )
327
- rope_parameters["mscale"] = 0 # Disable mscale
328
- rope_parameters["mscale_all_dim"] = 0 # Disable mscale
329
- rope_parameters["is_deepseek_v4"] = True
330
- rope_parameters["rope_dim"] = self.rope_head_dim
331
- self.rotary_emb = get_rope(
332
- self.head_dim,
333
- max_position=self.max_position_embeddings,
334
- rope_parameters=rope_parameters,
335
- is_neox_style=False,
317
+ self.rotary_emb = build_deepseek_v4_rope(
318
+ config,
319
+ head_dim=self.head_dim,
320
+ rope_head_dim=self.rope_head_dim,
321
+ max_position_embeddings=self.max_position_embeddings,
322
+ compress_ratio=self.compress_ratio,
323
)
324
325
self.indexer = None
vllm/models/deepseek_v4/common/rope.py
new
+36
@@ -0,0 +1,36 @@
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
+ rope_parameters["rope_theta"] = (
19
+ config.compress_rope_theta if compress_ratio > 1 else config.rope_theta
20
+ )
21
+ if rope_parameters["rope_type"] != "default":
22
+ rope_parameters["rope_type"] = (
23
+ "deepseek_yarn"
24
+ if rope_parameters.get("apply_yarn_scaling", True)
25
+ else "deepseek_llama_scaling"
26
+ )
27
+ rope_parameters["mscale"] = 0 # Disable mscale
28
+ rope_parameters["mscale_all_dim"] = 0 # Disable mscale
29
+ rope_parameters["is_deepseek_v4"] = True
30
+ rope_parameters["rope_dim"] = rope_head_dim
31
+ return get_rope(
32
+ head_dim,
33
+ max_position=max_position_embeddings,
34
+ rope_parameters=rope_parameters,
35
+ is_neox_style=False,
36
+ )
vllm/models/deepseek_v4/nvidia/model.py
+7
-20
@@ -35,7 +35,6 @@ from vllm.model_executor.layers.linear import (
35
)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
from vllm.model_executor.layers.quantization import QuantizationConfig
38
-from vllm.model_executor.layers.rotary_embedding import get_rope
38
from vllm.model_executor.layers.vocab_parallel_embedding import (
39
ParallelLMHead,
40
VocabParallelEmbedding,
@@ -56,6 +55,7 @@ from vllm.models.deepseek_v4.attention import (
55
DeepseekV4Indexer,
56
DeepseekV4MLA,
57
)
58
+from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
59
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
60
from vllm.sequence import IntermediateTensors
61
@@ -697,25 +697,12 @@ class DeepseekV4Attention(nn.Module):
697
self.rope_parameters = config.rope_scaling
698
699
# Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it)
700
- rope_parameters = config.rope_parameters
701
- rope_parameters["rope_theta"] = (
702
- config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta
703
- )
704
- if config.rope_parameters["rope_type"] != "default":
705
- config.rope_parameters["rope_type"] = (
706
- "deepseek_yarn"
707
- if config.rope_parameters.get("apply_yarn_scaling", True)
708
- else "deepseek_llama_scaling"
709
- )
710
- rope_parameters["mscale"] = 0 # Disable mscale
711
- rope_parameters["mscale_all_dim"] = 0 # Disable mscale
712
- rope_parameters["is_deepseek_v4"] = True
713
- rope_parameters["rope_dim"] = self.rope_head_dim
714
- self.rotary_emb = get_rope(
715
- self.head_dim,
716
- max_position=self.max_position_embeddings,
717
- rope_parameters=rope_parameters,
718
- is_neox_style=False,
700
+ self.rotary_emb = build_deepseek_v4_rope(
701
+ config,
702
+ head_dim=self.head_dim,
703
+ rope_head_dim=self.rope_head_dim,
704
+ max_position_embeddings=self.max_position_embeddings,
705
+ compress_ratio=self.compress_ratio,
706
)
707
708
self.indexer = None