From 6d7d4da99e41c4ccc0d52d74e2bf36d1ff31034d Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Wed, 29 Apr 2026 18:08:55 +0800 Subject: [PATCH] [Bugfix] BailingMoeV2.5: rotate full qk_rope_head_dim in MLA RoPE (#41185) Signed-off-by: zjy0516 --- vllm/model_executor/models/bailing_moe_linear.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 55ea1bad44d..bcf1200bd1b 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -205,13 +205,19 @@ class BailingMoeV25MLAAttention(nn.Module): self.q_a_layernorm = None self.q_b_proj = None - rope_parameters = _build_rope_parameters(config) + rope_parameters = _build_rope_parameters(config) or {} + # MLA rotates the full qk_rope_head_dim, + # partial_rotary_factor is for the linear-attn head only. + rope_parameters = { + k: v for k, v in rope_parameters.items() if k != "partial_rotary_factor" + } + rope_parameters["rope_dim"] = self.qk_rope_head_dim max_position = getattr(config, "max_position_embeddings", 8192) self.rotary_emb = get_rope( head_size=self.qk_rope_head_dim, max_position=max_position, is_neox_style=False, - rope_parameters=rope_parameters or None, + rope_parameters=rope_parameters, ) # Build MLAModules for MultiHeadLatentAttentionWrapper