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] BailingMoeV2.5: rotate full qk_rope_head_dim in MLA RoPE (#41185)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Apr 29, 2026 at 18:08 UTC
6d7d4da99e41c4ccc0d52d74e2bf36d1ff31034d
1 file changed
+8
-2
vllm/model_executor/models/bailing_moe_linear.py
+8
-2
@@ -205,13 +205,19 @@ class BailingMoeV25MLAAttention(nn.Module):
205
self.q_a_layernorm = None
206
self.q_b_proj = None
207
208
- rope_parameters = _build_rope_parameters(config)
208
+ rope_parameters = _build_rope_parameters(config) or {}
209
+ # MLA rotates the full qk_rope_head_dim,
210
+ # partial_rotary_factor is for the linear-attn head only.
211
+ rope_parameters = {
212
+ k: v for k, v in rope_parameters.items() if k != "partial_rotary_factor"
213
+ }
214
+ rope_parameters["rope_dim"] = self.qk_rope_head_dim
215
max_position = getattr(config, "max_position_embeddings", 8192)
216
self.rotary_emb = get_rope(
217
head_size=self.qk_rope_head_dim,
218
max_position=max_position,
219
is_neox_style=False,
214
- rope_parameters=rope_parameters or None,
220
+ rope_parameters=rope_parameters,
221
)
222
223
# Build MLAModules for MultiHeadLatentAttentionWrapper