[bugfix] fix mla rope (#8462)

Author: Jintao-Huang

swift/megatron/init.py

@@ -876,9 +876,9 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor):
             # x_pe   [seqlen, batch, *, qk_pos_emb_head_dim]
             x_pe, x_nope = torch.split(
                 x, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1)
-            origin_rotary_interleaved = self.config.rotary_interleaved
+            origin_multi_latent_attention = self.config.multi_latent_attention
             try:
-                self.config.rotary_interleaved = self.config.dsa_indexer_rotary_interleaved
+                self.config.multi_latent_attention = self.config.dsa_indexer_rotary_interleaved
                 x_pe = apply_rotary_pos_emb(
                     x_pe,
                     rotary_pos_emb,
@@ -887,7 +887,7 @@ def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor):
                     cp_group=self.pg_collection.cp,
                 )
             finally:
-                self.config.rotary_interleaved = origin_rotary_interleaved
+                self.config.multi_latent_attention = origin_multi_latent_attention
             # [seqlen, batch, *, index_head_dim]
             x = torch.cat([x_pe, x_nope], dim=-1)
             return x

swift/megatron/model/gpt_model.py

@@ -176,6 +176,10 @@ def _apply_rotary_pos_emb_bshd(
 
             # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
             t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+            if multi_latent_attention:
+                x1 = t[..., 0::2]
+                x2 = t[..., 1::2]
+                t = torch.cat((x1, x2), dim=-1)
 
             # first part is cosine component
             # second part is sine component, need to change signs with _rotate_half method

swift/megatron/model/model_config.py

@@ -278,11 +278,7 @@ def __post_init__(self):
             assert not self.swiglu
             self.gated_linear_unit = True
             self.activation_func = quick_gelu
-        _origin_rotary_interleaved = self.rotary_interleaved
-        if self.multi_latent_attention and self.rotary_interleaved:
-            self.rotary_interleaved = False
         super().__post_init__()
-        self.rotary_interleaved = _origin_rotary_interleaved
         self._check_npu()
         self.variable_seq_lengths = True
 
@@ -481,8 +477,6 @@ def convert_hf_config(config) -> Dict[str, Any]:
             res.pop('num_query_groups', None)
         if llm_model_type == 'glm_moe_dsa':
             res['experimental_attention_variant'] = 'dsa'
-            # https://github.com/modelscope/ms-swift/pull/8085
-            # res['rotary_interleaved'] = False
     elif llm_model_type == 'qwen3_next' or hf_model_type in {'qwen3_5', 'qwen3_5_moe'}:
         use_mcore_gdn = get_env_args('SWIFT_USE_MCORE_GDN', bool, False)
         if use_mcore_gdn and llm_model_type == 'qwen3_next':
@@ -525,10 +519,6 @@ def convert_hf_config(config) -> Dict[str, Any]:
         mrope_interleaved = rope_scaling.get('mrope_interleaved', False) or rope_scaling.get('interleaved', False)
         res['mrope_interleaved'] = mrope_interleaved
 
-    if res.get('multi_latent_attention') and res.get('position_embedding_type') in {
-            'rope', None
-    } and 'rotary_interleaved' not in res:
-        res['rotary_interleaved'] = True
     if first_k_dense_replace is not None:
         res['moe_layer_freq'] = f'[0]*{first_k_dense_replace}+[1]*{res["num_layers"] - first_k_dense_replace}'
     if res.get('moe_router_score_function', 'softmax') == 'sigmoid' and 'moe_router_enable_expert_bias' not in res: