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 11, 00:00 UTC
[Bug] Fix `Number of dimensions of tensors must match.` for Deepseek V3.2 (#31160)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 23, 2025 at 21:41 UTC
76e6a951925bf37c49f88ad155dc9fcec01a3faf
1 file changed
+6
-3
vllm/model_executor/models/deepseek_v2.py
+6
-3
@@ -878,11 +878,14 @@ class Indexer(nn.Module):
878
)
879
880
q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
881
- # `rotary_emb` is shape-preserving; `q_pe` is already
882
- # [num_tokens, n_head, rope_dim].
881
+ # Note: RoPE (NeoX) can introduce extra leading dimensions during compilation
882
+ # so we need to reshape back to token-flattened shapes
883
+ q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
884
+ k_pe = k_pe.reshape(-1, 1, self.rope_dim)
885
+
886
q = torch.cat([q_pe, q_nope], dim=-1)
887
# `k_pe` is [num_tokens, 1, rope_dim] (MQA).
885
- k = torch.cat([k_pe.squeeze(1), k_nope], dim=-1)
888
+ k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
889
890
# we only quant q here since k quant is fused with cache insertion
891
q = q.view(-1, self.head_dim)