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 4, 12:30 UTC
[Bugfix] Fix repeated DSv4 RoPE cache initialization (#41148)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Jee Jee Li committed
Apr 29, 2026 at 20:29 UTC
9d8ad5b408bf447e41a3629fc21a453720aaf52b
2 files changed
+11
-3
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+11
-2
@@ -45,6 +45,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
45
beta_slow: int = 1,
46
mscale: float = 1,
47
mscale_all_dim: float = 0,
48
+ init_cache: bool = True,
49
) -> None:
50
self.scaling_factor = scaling_factor
51
self.extrapolation_factor = extrapolation_factor
@@ -65,7 +66,13 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
66
and head_size in [64, 128, 256, 512]
67
)
68
super().__init__(
68
- head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
69
+ head_size,
70
+ rotary_dim,
71
+ max_position_embeddings,
72
+ base,
73
+ is_neox_style,
74
+ dtype,
75
+ init_cache=init_cache,
76
)
77
78
def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
@@ -211,7 +218,9 @@ class DeepseekV4ScalingRotaryEmbedding(DeepseekScalingRotaryEmbedding):
218
"""
219
220
def __init__(self, *args, **kwargs):
214
- super().__init__(*args, **kwargs)
221
+ # Avoid compute cache repeatedly
222
+ kwargs.pop("init_cache", None)
223
+ super().__init__(*args, **kwargs, init_cache=False)
224
cache_fp32 = self._compute_cos_sin_cache()
225
self.register_buffer("cos_sin_cache", cache_fp32, persistent=False)
226
vllm/model_executor/models/deepseek_v4.py
-1
@@ -1027,7 +1027,6 @@ class DeepseekV4Attention(nn.Module):
1027
max_position=self.max_position_embeddings,
1028
rope_parameters=rope_parameters,
1029
is_neox_style=False,
1030
- dtype=config.torch_dtype,
1030
)
1031
1032
self.indexer = None