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 12, 22:15 UTC
[DeepSeek] Fix DeepSeek V3.2 Rope Embedding (#28968)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Nov 19, 2025 at 16:30 UTC
88f5b19f0bc681c016eaaa17502d3bb4e2b59b51
2 files changed
+17
-3
vllm/model_executor/layers/mla.py
+5
-1
@@ -24,6 +24,7 @@ class MLAModules:
24
q_b_proj: torch.nn.Module | None
25
q_proj: torch.nn.Module | None
26
indexer: torch.nn.Module | None
27
+ indexer_rotary_emb: torch.nn.Module | None
28
is_sparse: bool
29
topk_indices_buffer: torch.Tensor | None
30
@@ -80,6 +81,7 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
81
self.rotary_emb = mla_modules.rotary_emb
82
self.o_proj = mla_modules.o_proj
83
self.indexer = mla_modules.indexer
84
+ self.indexer_rope_emb = mla_modules.indexer_rotary_emb
85
self.is_sparse = mla_modules.is_sparse
86
87
if self.indexer is not None:
@@ -153,7 +155,9 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
155
)
156
157
if self.indexer and self.is_sparse:
156
- _topk_indices = self.indexer(hidden_states, q_c, positions, self.rotary_emb)
158
+ _topk_indices = self.indexer(
159
+ hidden_states, q_c, positions, self.indexer_rope_emb
160
+ )
161
162
attn_out = self.mla_attn(
163
q,
vllm/model_executor/models/deepseek_v2.py
+12
-2
@@ -837,8 +837,8 @@ class Indexer(nn.Module):
837
)
838
839
q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
840
- q = torch.cat([q_pe, q_nope], dim=-1)
841
- k = torch.cat([k_pe.squeeze(1), k_nope], dim=-1)
840
+ q = torch.cat([q_pe.squeeze(0), q_nope], dim=-1)
841
+ k = torch.cat([k_pe.squeeze((0, 2)), k_nope], dim=-1)
842
843
# we only quant q here since k quant is fused with cache insertion
844
q = q.view(-1, self.head_dim)
@@ -987,6 +987,14 @@ class DeepseekV2MLAAttention(nn.Module):
987
self.is_v32 = hasattr(config, "index_topk")
988
989
if self.is_v32:
990
+ self.indexer_rope_emb = get_rope(
991
+ qk_rope_head_dim,
992
+ rotary_dim=qk_rope_head_dim,
993
+ max_position=max_position_embeddings,
994
+ base=rope_theta,
995
+ rope_scaling=rope_scaling,
996
+ is_neox_style=True,
997
+ )
998
self.indexer = Indexer(
999
vllm_config,
1000
config,
@@ -998,6 +1006,7 @@ class DeepseekV2MLAAttention(nn.Module):
1006
f"{prefix}.indexer",
1007
)
1008
else:
1009
+ self.indexer_rope_emb = None
1010
self.indexer = None
1011
1012
mla_modules = MLAModules(
@@ -1015,6 +1024,7 @@ class DeepseekV2MLAAttention(nn.Module):
1024
q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
1025
q_proj=self.q_proj if self.q_lora_rank is None else None,
1026
indexer=self.indexer,
1027
+ indexer_rotary_emb=self.indexer_rope_emb,
1028
is_sparse=self.is_v32,
1029
topk_indices_buffer=topk_indices_buffer,
1030
)