From eceae54e45e3a7c6068cb86318986924940f39d3 Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Thu, 9 Jul 2026 17:56:39 +0800
Subject: [PATCH 1/7] support absorbed_mla

---
 src/mcore_bridge/model/modules/__init__.py    |   1 +
 .../model/modules/absorbed_mla.py             | 338 ++++++++++++++++++
 src/mcore_bridge/model/register.py            |   6 +-
 3 files changed, 343 insertions(+), 2 deletions(-)
 create mode 100644 src/mcore_bridge/model/modules/absorbed_mla.py

diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py
index edf406d..7269abd 100644
--- a/src/mcore_bridge/model/modules/__init__.py
+++ b/src/mcore_bridge/model/modules/__init__.py
@@ -1,4 +1,5 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+from .absorbed_mla import AbsorbedMLASelfAttention
 from .compressor import Compressor, CSAIndexer
 from .dsa_indexer import DSAIndexer
 from .gated_delta_net import GatedDeltaNet
diff --git a/src/mcore_bridge/model/modules/absorbed_mla.py b/src/mcore_bridge/model/modules/absorbed_mla.py
new file mode 100644
index 0000000..77bf42e
--- /dev/null
+++ b/src/mcore_bridge/model/modules/absorbed_mla.py
@@ -0,0 +1,338 @@
+try:
+    from megatron.core.transformer.experimental_attention_variant.absorbed_mla import \
+        AbsorbedMLASelfAttention as McoreAbsorbedMLASelfAttention
+except ImportError:
+    McoreAbsorbedMLASelfAttention = None
+
+import torch
+from megatron.core import tensor_parallel
+from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
+from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region,
+                                                    gather_from_tensor_model_parallel_region,
+                                                    scatter_to_sequence_parallel_region)
+from megatron.core.utils import deprecate_inference_params
+
+
+class AbsorbedMLASelfAttention(McoreAbsorbedMLASelfAttention):
+
+    def get_query_key_value_tensors(
+        self,
+        hidden_states,
+        key_value_states=None,
+        packed_seq_params=None,
+        inference_context=None,
+        *,
+        inference_params=None,
+    ):
+        """
+        Derives absorbed q, compressed q, and compressed kv tensors from `hidden_states`.
+        """
+        # s = sequence length, b = batch size, h = hidden size
+        from megatron.core.utils import get_pg_size
+        assert (hidden_states.ndim == 3), f"hidden_states should be 3D, [s, b, h], got {hidden_states.ndim}D"
+        if packed_seq_params is not None:
+            assert (packed_seq_params.local_cp_size
+                    is None), 'dynamic context parallel is not supported with MLA yet and is planned for future. \
+            Please disable dynamic context parallel.'
+
+        inference_context = deprecate_inference_params(inference_context, inference_params)
+
+        # =========================================
+        # Prepare RoPE and seqlen related params
+        # =========================================
+        rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(inference_context, None, hidden_states, self.config,
+                                                                packed_seq_params)
+
+        mscale = 1.0
+        packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'
+        if self.config.rope_type == 'rope':
+            rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
+        else:
+            rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
+
+        if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
+            if packed_seq_params.cu_seqlens_q_padded is not None:
+                cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded
+            else:
+                cu_seqlens_q = packed_seq_params.cu_seqlens_q
+            if packed_seq_params.cu_seqlens_kv_padded is not None:
+                cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded
+            else:
+                cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+        else:
+            cu_seqlens_q = cu_seqlens_kv = None
+
+        # =========================================
+        # Q down projection
+        # =========================================
+        if self.config.q_lora_rank is not None:
+            # if linear_q_down_proj is ColumnParallelLinear:
+            #     q_compressed: [s, b, q_lora_rank / TP]
+            # elif linear_q_down_proj is Linear:
+            #     q_compressed: [s / TP, b, q_lora_rank]
+            q_compressed, _ = self.linear_q_down_proj(hidden_states)
+
+            # When output is sharded (ColumnParallelLinear), two things are needed to be
+            # identical to a normal Linear.
+            #   1. Manually gather output to restore output dim q_lora_rank;
+            #   2. Scatter sequence back to s / TP if sequence-parallel since it was
+            #      gathered by ColumnParallelLinear.
+            if q_compressed.size(-1) != self.config.q_lora_rank:
+                q_compressed = gather_from_tensor_model_parallel_region(q_compressed)
+                if self.config.sequence_parallel:
+                    q_compressed = scatter_to_sequence_parallel_region(q_compressed)
+        else:
+            q_compressed = hidden_states
+
+        # =========================================
+        # KV down projection
+        # =========================================
+        # if linear_kv_down_proj is ColumnParallelLinear:
+        #     kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim) / TP]
+        # elif linear_kv_down_proj is Linear:
+        #     kv_combined: [s / TP, b, (kv_lora_rank + qk_pos_emb_head_dim)]
+        kv_combined, _ = self.linear_kv_down_proj(hidden_states)
+        if kv_combined.size(-1) != self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim:
+            # kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim)]
+            kv_combined = gather_from_tensor_model_parallel_region(kv_combined)
+            # kv_compressed:[s, b, kv_lora_rank], k_pos_emb: [s, b, qk_pos_emb_head_dim]
+            kv_compressed, k_pos_emb = torch.split(
+                kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1)
+            if self.config.sequence_parallel:
+                # kv_compressed:[s / TP, b, kv_lora_rank]
+                kv_compressed = scatter_to_sequence_parallel_region(kv_compressed)
+        else:
+            # kv_compressed:[s / TP, b, kv_lora_rank], k_pos_emb: [s / TP, b, qk_pos_emb_head_dim]
+            kv_compressed, k_pos_emb = torch.split(
+                kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1)
+            if get_pg_size(self.tp_group) > 1 and self.config.sequence_parallel:
+                # k_pos_emb: [s, b, qk_pos_emb_head_dim]
+                k_pos_emb = gather_from_sequence_parallel_region(k_pos_emb, group=self.tp_group)
+
+        if packed_seq_params is not None:
+            assert q_compressed.ndim == 3 and q_compressed.size(1) == 1
+            assert kv_compressed.ndim == 3 and kv_compressed.size(1) == 1
+            assert k_pos_emb.ndim == 3 and k_pos_emb.size(1) == 1
+            # If sequence packing, TE expect [t, h, d] shaped qkv input.
+            # In Megatron-Core, the qkv shape is [t, 1, h, d].
+            # So we need to reshape qkv from [t, 1, h, d] to [t, h, d].
+            q_compressed = q_compressed.squeeze(1)
+            kv_compressed = kv_compressed.squeeze(1)
+            k_pos_emb = k_pos_emb.squeeze(1)
+
+        # =========================================
+        # Apply norm
+        # =========================================
+        if self.config.q_lora_rank is not None:
+            # q_compressed: [num_tokens, q_lora_rank]
+            q_compressed = self.q_layernorm(q_compressed)
+
+        kv_compressed = self.kv_layernorm(kv_compressed)
+        # Because we won't apply V up projection to the compressed KV, so we need to gather it
+        # manually.
+        if get_pg_size(self.tp_group) > 1 and self.config.sequence_parallel:
+            kv_compressed = gather_from_sequence_parallel_region(kv_compressed, group=self.tp_group)
+
+        # =========================================
+        # QKV up projection and RoPE apply
+        # =========================================
+
+        def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb):
+            """
+            Apply the up projection and RoPE to the query and key.
+            When sequence packing enabled, the input tensors adopt a packed shape of [t, ...];
+            otherwise, they maintain the unpacked shape [s, b, ...]. In subsequent code comments,
+            we uniformly use [num_tokens, ...] to denote [s, b, ...] or [t, ...] for two cases.
+            """
+            if self.config.q_lora_rank is not None:
+                # q_compressed: [num_tokens, q_lora_rank]
+                # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)]
+                q, _ = self.linear_q_up_proj(q_compressed)
+            else:
+                # q_compressed: [num_tokens, hidden_size]
+                # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)]
+                q, _ = self.linear_q_proj(q_compressed)
+
+            # q: [num_tokens, n, q_head_dim]
+            q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim)
+
+            # [num_tokens, kv_lora_rank] -> [num_tokens, 1, kv_lora_rank]
+            kv_compressed = torch.unsqueeze(kv_compressed, -2)
+            # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim]
+            k_pos_emb = torch.unsqueeze(k_pos_emb, -2)
+
+            k_up_weight, _ = self._get_kv_up_weights()
+
+            q_len = q.size()[0]
+            if inference_context is not None:
+                # add offset to the sequence start for inference
+                sequence_start = inference_context.sequence_len_offset
+                sequence_end = sequence_start + q_len
+                rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end]
+            elif packed_seq_params is None or self.config.context_parallel_size == 1:
+                # Shorten rotary_pos_emb to the sequence length when inference_params
+                # is not provided. This makes sure we can run forward directly with
+                # any sequence length. During training, the sequence length is always
+                # the full rotary_pos_emb length, except for sequence packing + CP.
+                # When sequence packing and context parallel are both enabled, the
+                # position embedding will not split rotary_pos_emb, so it may exceed
+                # the sequence length on this CP rank, but we need the full rotary_pos_emb
+                # to cover the full sequence, so we do not shorten it here.
+                rotary_pos_emb = rotary_pos_emb[0:q_len]
+
+            # q_no_pe: [num_tokens, n, qk_head_dim]
+            # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim]
+            q_no_pe, q_pos_emb = torch.split(q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1)
+
+            # Absorb k_up_weight into q_no_pe
+            # q_absorbed: [num_tokens, n, kv_lora_rank]
+            q_absorbed = torch.einsum('...nd,ndk->...nk', q_no_pe, k_up_weight)
+            q_absorbed = q_absorbed.contiguous()
+            assert q_absorbed.ndim == q.ndim
+            assert q_absorbed.shape[:-1] == q.shape[:-1]
+            assert q_absorbed.size(-1) == self.config.kv_lora_rank
+
+            # Apply RoPE to q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim]
+            q_pos_emb = apply_rotary_pos_emb(
+                q_pos_emb,
+                rotary_pos_emb,
+                config=self.config,
+                cu_seqlens=cu_seqlens_q,
+                mscale=mscale,
+                cp_group=self.pg_collection.cp,
+                mla_rotary_interleaved=True,
+            )
+            # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim]
+            k_pos_emb = apply_rotary_pos_emb(
+                k_pos_emb,
+                rotary_pos_emb,
+                config=self.config,
+                cu_seqlens=cu_seqlens_kv,
+                mscale=mscale,
+                cp_group=self.pg_collection.cp,
+                mla_rotary_interleaved=True,
+            )
+
+            # query: [num_tokens, n, (kv_lora_rank + qk_pos_emb_head_dim)]
+            q_absorbed = torch.cat([q_absorbed, q_pos_emb], dim=-1)
+            # key: [num_tokens, 1, (kv_lora_rank + qk_pos_emb_head_dim)]
+            kv_compressed = torch.cat([kv_compressed, k_pos_emb], dim=-1)
+
+            assert q_absorbed.is_contiguous()
+            assert kv_compressed.is_contiguous()
+
+            return q_absorbed, kv_compressed
+
+        if self.recompute_up_proj:
+            quantization = self.config.fp8 or self.config.fp4
+            assert not quantization, 'FP8/FP4 is not supported for AbsorbedMLA'
+            self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization)
+            q_absorbed, kv_compressed = self.qkv_up_checkpoint.checkpoint(qkv_up_proj_and_rope_apply, q_compressed,
+                                                                          kv_compressed, k_pos_emb, rotary_pos_emb)
+        else:
+            assert not self.cache_mla_latents, 'cache_mla_latents is not supported for AbsorbedMLA'
+            q_absorbed, kv_compressed = qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb,
+                                                                   rotary_pos_emb)
+
+        return q_absorbed, kv_compressed, q_compressed
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        key_value_states=None,
+        inference_context=None,
+        rotary_pos_emb=None,
+        rotary_pos_cos=None,
+        rotary_pos_sin=None,
+        rotary_pos_cos_sin=None,
+        attention_bias=None,
+        packed_seq_params=None,
+        position_ids=None,
+        sequence_len_offset=None,
+        *,
+        inference_params=None,
+    ):
+        from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
+            _apply_absorbed_v_up_projection, _restore_packed_thd_batch_dim)
+        """Forward pass for multi-latent attention with matrix absorption"""
+        assert rotary_pos_emb is None, 'Rotary position embeddings should not be passed into MLA.'
+        assert attention_bias is None, 'Attention bias should not be passed into MLA.'
+        assert (rotary_pos_cos is None and rotary_pos_sin is None), 'MLA does not support Flash Decoding'
+        assert not rotary_pos_cos_sin, 'Flash-infer rope has not been tested with MLA.'
+        assert not (self.training and self.cache_mla_latents), 'cache_mla_latents conflicts with training.'
+        assert (inference_context is None and inference_params is None), 'Inference is not supported for AbsorbedMLA'
+
+        # =====================
+        # Query, Key, and Value
+        # =====================
+        q_absorbed, kv_compressed, q_compressed = self.get_query_key_value_tensors(
+            hidden_states, key_value_states, packed_seq_params, inference_context=inference_context)
+
+        assert q_absorbed.is_contiguous()
+        assert q_compressed.is_contiguous()
+        assert kv_compressed.is_contiguous()
+        v_up_weight = self._get_v_up_weight()
+
+        # ==================================
+        # Core attention computation
+        # ==================================
+        if self.checkpoint_core_attention and self.training:
+            core_attn_out = self._checkpointed_attention_forward(
+                q_absorbed,
+                kv_compressed,
+                hidden_states,
+                q_compressed,
+                attention_mask,
+                v_up_weight,
+                position_ids=position_ids,
+                packed_seq_params=packed_seq_params,
+            )
+        else:
+            core_attn_out = self.core_attention(
+                q_absorbed,
+                kv_compressed,
+                value=None,
+                attention_mask=attention_mask,
+                x=hidden_states,
+                qr=q_compressed,
+                up_v_weight=v_up_weight,
+                position_ids=position_ids,
+                packed_seq_params=packed_seq_params,
+                attn_mask_type=self.attn_mask_type,
+            )
+
+        # ==================================
+        # Apply V up projection
+        # ==================================
+        core_consumed_v_up_projection = getattr(self.core_attention, 'consumes_absorbed_v_up_projection', False)
+        core_attn_out = _apply_absorbed_v_up_projection(
+            core_attn_out,
+            v_up_weight,
+            self.num_attention_heads_per_partition,
+            self.config.kv_lora_rank,
+            self.config.v_head_dim,
+            core_consumed_v_up_projection,
+        )
+
+        core_attn_out = _restore_packed_thd_batch_dim(core_attn_out, hidden_states, packed_seq_params)
+
+        assert core_attn_out.ndim == hidden_states.ndim
+        assert core_attn_out.shape[0] == (hidden_states.shape[0] * self.config.tensor_model_parallel_size), (
+            f"{core_attn_out.shape[0]} != "
+            f"{hidden_states.shape[0]} * "
+            f"{self.config.tensor_model_parallel_size}")
+        assert core_attn_out.shape[1:-1] == hidden_states.shape[1:-1]
+        assert core_attn_out.size(-1) == (self.config.v_head_dim * self.num_attention_heads_per_partition)
+
+        if self.recompute_up_proj:
+            assert self.qkv_up_checkpoint is not None
+            self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out)
+            self.qkv_up_checkpoint = None
+
+        # =================
+        # Output. [sq, b, h]
+        # =================
+        output, bias = self.linear_proj(core_attn_out)
+
+        return output, bias
diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py
index 5f39664..4bf1a7a 100644
--- a/src/mcore_bridge/model/register.py
+++ b/src/mcore_bridge/model/register.py
@@ -20,8 +20,8 @@
 from mcore_bridge.config import ModelConfig
 from mcore_bridge.utils import get_logger
 
-from .modules import (DSAIndexer, MLASelfAttention, MultiTokenPredictionLayer, TopKRouter, TransformerBlock,
-                      TransformerLayer)
+from .modules import (AbsorbedMLASelfAttention, DSAIndexer, MLASelfAttention, MultiTokenPredictionLayer, TopKRouter,
+                      TransformerBlock, TransformerLayer)
 
 if TYPE_CHECKING:
     from .gpt_model import GPTModel
@@ -163,6 +163,8 @@ def _replace_mla_attention(self, transformer_layer_spec):
             self_attention = layer_spec.submodules.self_attention
             if self_attention.module is McoreMLASelfAttention:
                 self_attention.module = MLASelfAttention
+            elif self_attention.module.__name__ == 'AbsorbedMLASelfAttention':
+                self_attention.module = AbsorbedMLASelfAttention
 
     def _replace_router(self, transformer_layer_spec, mlp_key='mlp'):
         for layer_spec in transformer_layer_spec.layer_specs:

From e8e6e7f57b278de8e7227eea13c9877b4618c35e Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 14 Jul 2026 11:10:28 +0800
Subject: [PATCH 2/7] update

---
 src/mcore_bridge/config/model_config.py       |  2 ++
 src/mcore_bridge/config/parser.py             |  2 ++
 src/mcore_bridge/model/gpts/glm_moe_dsa.py    | 22 +++++++++++-----
 .../model/modules/absorbed_mla.py             | 26 +++++++++----------
 src/mcore_bridge/model/modules/dsa_indexer.py | 21 ++++++++++-----
 5 files changed, 47 insertions(+), 26 deletions(-)

diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py
index ce07979..e6d0294 100644
--- a/src/mcore_bridge/config/model_config.py
+++ b/src/mcore_bridge/config/model_config.py
@@ -204,6 +204,8 @@ class ModelConfig(TransformerConfig):
     dsa_indexer_loss_coeff: float = 0.
     dsa_indexer_use_sparse_loss: bool = False
     dsa_indexer_rotary_interleaved: bool = False
+    dsa_indexer_topk_freq: int = 1
+    dsa_indexer_skip_topk_offset: int = 0
 
     # deepseek-v4
     csa_window_size: int = 128
diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index 7cf59c4..54b8e05 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -56,6 +56,8 @@
     'dsa_indexer_head_dim': ['index_head_dim'],
     'dsa_indexer_topk': ['index_topk'],
     'dsa_indexer_rotary_interleaved': ['indexer_rope_interleave'],
+    'dsa_indexer_topk_freq': ['index_topk_freq'],
+    'dsa_indexer_skip_topk_offset': ['index_skip_topk_offset'],
     # deepseek_v4
     'csa_compress_ratios': ['compress_rates'],
     'csa_compress_rotary_base': ['compress_rope_theta'],
diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
index 9d4054c..b73f2d2 100644
--- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py
+++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
@@ -12,14 +12,24 @@
 
 try:
     from megatron.core.transformer.experimental_attention_variant.dsa import (DSAIndexerLossAutoScaler,
-                                                                              DSAIndexerLossLoggingHelper, DSAttention,
-                                                                              FusedDSAIndexerLoss, unfused_dsa_fn)
+                                                                              DSAIndexerLossLoggingHelper)
+    from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention as McoreDSAttention
+    from megatron.core.transformer.experimental_attention_variant.dsa import FusedDSAIndexerLoss, unfused_dsa_fn
 except ImportError:
-    DSAttention = object
+    McoreDSAttention = object
 
 mcore_019 = version.parse(megatron.core.__version__) >= version.parse('0.19.0rc0')
 
 
+class DSAttention(McoreDSAttention):
+
+    def _get_index_share_carrier(self, packed_seq_params, attention_mask):
+        """Return the object that carries DSA top-k sharing state for this forward."""
+        if packed_seq_params is not None and packed_seq_params.qkv_format is not None:
+            return packed_seq_params
+        return attention_mask if attention_mask is not None else self.config
+
+
 class GlmMoeDsaDSAttention(DSAttention):
     """DSAttention with shared indexer support for GLM 5.2.
 
@@ -170,8 +180,8 @@ def _layer_forward(self, layer, hidden_states, **kwargs):
 
 
 class GlmMoeDsaLoader(ModelLoader):
-    model_cls = GlmMoeDsaGPTModel
-    transformer_block = GlmMoeDsaTransformerBlock
+    model_cls = GPTModel if mcore_019 else GlmMoeDsaGPTModel
+    transformer_block = TransformerBlock if mcore_019 else GlmMoeDsaTransformerBlock
 
     def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
         transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
@@ -181,7 +191,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
             for layer_spec in transformer_layer_spec.layer_specs:
                 core_attn = layer_spec.submodules.self_attention.submodules.core_attention
                 if hasattr(core_attn, 'module') and issubclass(core_attn.module, DSAttention):
-                    core_attn.module = GlmMoeDsaDSAttention
+                    core_attn.module = DSAttention if mcore_019 else GlmMoeDsaDSAttention
 
         return transformer_layer_spec
 
diff --git a/src/mcore_bridge/model/modules/absorbed_mla.py b/src/mcore_bridge/model/modules/absorbed_mla.py
index 77bf42e..0a87a41 100644
--- a/src/mcore_bridge/model/modules/absorbed_mla.py
+++ b/src/mcore_bridge/model/modules/absorbed_mla.py
@@ -7,6 +7,7 @@
 import torch
 from megatron.core import tensor_parallel
 from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
+from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region,
                                                     gather_from_tensor_model_parallel_region,
                                                     scatter_to_sequence_parallel_region)
@@ -21,6 +22,7 @@ def get_query_key_value_tensors(
         key_value_states=None,
         packed_seq_params=None,
         inference_context=None,
+        rotary_pos_emb=None,
         *,
         inference_params=None,
     ):
@@ -40,16 +42,6 @@ def get_query_key_value_tensors(
         # =========================================
         # Prepare RoPE and seqlen related params
         # =========================================
-        rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(inference_context, None, hidden_states, self.config,
-                                                                packed_seq_params)
-
-        mscale = 1.0
-        packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'
-        if self.config.rope_type == 'rope':
-            rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
-        else:
-            rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)
-
         if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd':
             if packed_seq_params.cu_seqlens_q_padded is not None:
                 cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded
@@ -198,7 +190,6 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
                 rotary_pos_emb,
                 config=self.config,
                 cu_seqlens=cu_seqlens_q,
-                mscale=mscale,
                 cp_group=self.pg_collection.cp,
                 mla_rotary_interleaved=True,
             )
@@ -208,7 +199,6 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
                 rotary_pos_emb,
                 config=self.config,
                 cu_seqlens=cu_seqlens_kv,
-                mscale=mscale,
                 cp_group=self.pg_collection.cp,
                 mla_rotary_interleaved=True,
             )
@@ -256,7 +246,6 @@ def forward(
         from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
             _apply_absorbed_v_up_projection, _restore_packed_thd_batch_dim)
         """Forward pass for multi-latent attention with matrix absorption"""
-        assert rotary_pos_emb is None, 'Rotary position embeddings should not be passed into MLA.'
         assert attention_bias is None, 'Attention bias should not be passed into MLA.'
         assert (rotary_pos_cos is None and rotary_pos_sin is None), 'MLA does not support Flash Decoding'
         assert not rotary_pos_cos_sin, 'Flash-infer rope has not been tested with MLA.'
@@ -267,7 +256,11 @@ def forward(
         # Query, Key, and Value
         # =====================
         q_absorbed, kv_compressed, q_compressed = self.get_query_key_value_tensors(
-            hidden_states, key_value_states, packed_seq_params, inference_context=inference_context)
+            hidden_states,
+            key_value_states,
+            packed_seq_params,
+            rotary_pos_emb=rotary_pos_emb,
+            inference_context=inference_context)
 
         assert q_absorbed.is_contiguous()
         assert q_compressed.is_contiguous()
@@ -277,6 +270,11 @@ def forward(
         # ==================================
         # Core attention computation
         # ==================================
+        if self.config.experimental_attention_variant == 'dsa':
+            if packed_seq_params is None:
+                packed_seq_params = PackedSeqParams()
+            # for easy injection of rotary_pos_emb (patch)
+            packed_seq_params.rotary_pos_emb = rotary_pos_emb
         if self.checkpoint_core_attention and self.training:
             core_attn_out = self._checkpointed_attention_forward(
                 q_absorbed,
diff --git a/src/mcore_bridge/model/modules/dsa_indexer.py b/src/mcore_bridge/model/modules/dsa_indexer.py
index e148a53..d069e1c 100644
--- a/src/mcore_bridge/model/modules/dsa_indexer.py
+++ b/src/mcore_bridge/model/modules/dsa_indexer.py
@@ -42,8 +42,11 @@ def forward_before_topk(
         # =========================================
         # Gather inputs if sp is enabled
         # =========================================
-        packed_seq_params, rotary_pos_emb = packed_seq_params  # patch
-        assert packed_seq_params is None, 'Packed sequence is not supported for DSAttention'
+        if isinstance(packed_seq_params, tuple):
+            packed_seq_params, rotary_pos_emb = packed_seq_params  # patch
+        elif packed_seq_params is not None:
+            rotary_pos_emb = packed_seq_params.rotary_pos_emb
+        cu_seqlens = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None
 
         if self.config.sequence_parallel and self.pg_collection.tp.size() > 1:
             x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp)
@@ -62,7 +65,7 @@ def forward_before_topk(
         # [seqlen, batch, index_n_heads * index_head_dim]
         #   -> [seqlen, batch, index_n_heads, index_head_dim]
         q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim)
-        q = self._apply_rope(q, rotary_pos_emb)  # mscale will be passed in by patch
+        q = self._apply_rope(q, rotary_pos_emb, cu_seqlens)  # mscale will be passed in by patch
 
         # =========================================
         # k linear and apply rope to k
@@ -72,7 +75,7 @@ def forward_before_topk(
         k = self.k_norm(k)
         # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim]
         k = k.reshape(seqlen, bsz, 1, self.index_head_dim)
-        k = self._apply_rope(k, rotary_pos_emb)
+        k = self._apply_rope(k, rotary_pos_emb, cu_seqlens)
         # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim]
         k = k.reshape(seqlen, bsz, self.index_head_dim)
 
@@ -91,24 +94,30 @@ def forward_before_topk(
 
         return q, k, weights
 
-    def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor):
+    def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, cu_seqlens: Optional[torch.Tensor] = None):
         """Apply RoPE to the input tensor."""
         # x_nope [seqlen, batch, *, index_head_dim - qk_pos_emb_head_dim]
         # 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_multi_latent_attention = self.config.multi_latent_attention
+        squeezed_batch_dim = False
+        if cu_seqlens is not None and x_pe.ndim == 4 and x_pe.size(1) == 1:
+            x_pe = x_pe.squeeze(1)
+            squeezed_batch_dim = True
         try:
             self.config.multi_latent_attention = self.config.dsa_indexer_rotary_interleaved
             x_pe = apply_rotary_pos_emb(
                 x_pe,
                 rotary_pos_emb,
                 config=self.config,
-                cu_seqlens=None,
+                cu_seqlens=cu_seqlens,
                 cp_group=self.pg_collection.cp,
             )
         finally:
             self.config.multi_latent_attention = origin_multi_latent_attention
+        if squeezed_batch_dim:
+            x_pe = x_pe.unsqueeze(1)
         # [seqlen, batch, *, index_head_dim]
         x = torch.cat([x_pe, x_nope], dim=-1)
         return x

From 177cf6555ddcf3abe8fce561a6eb1fd50679d2b9 Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 14 Jul 2026 14:39:30 +0800
Subject: [PATCH 3/7] update

---
 src/mcore_bridge/model/gpts/glm_moe_dsa.py    | 162 +-----------------
 .../model/modules/absorbed_mla.py             |  10 ++
 2 files changed, 15 insertions(+), 157 deletions(-)

diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
index b73f2d2..b98b1a0 100644
--- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py
+++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
@@ -1,13 +1,9 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import megatron.core
-import torch
-from contextlib import contextmanager
 from packaging import version
 from typing import Optional
 
 from ..constant import ModelType
-from ..gpt_model import GPTModel
-from ..modules import TransformerBlock
 from ..register import ModelLoader, ModelMeta, register_model
 
 try:
@@ -30,168 +26,20 @@ def _get_index_share_carrier(self, packed_seq_params, attention_mask):
         return attention_mask if attention_mask is not None else self.config
 
 
-class GlmMoeDsaDSAttention(DSAttention):
-    """DSAttention with shared indexer support for GLM 5.2.
-
-    Refer: https://arxiv.org/abs/2603.12201 for more details.
-    """
-
-    def __init__(self, config, submodules, layer_number, *args, **kwargs):
-        super().__init__(config, submodules, layer_number, *args, **kwargs)
-        indexer_types = config.hf_config.indexer_types
-        self.skip_topk = False
-        if indexer_types is not None:
-            layer_idx = layer_number - 1
-            if layer_idx < len(indexer_types):
-                self.skip_topk = indexer_types[layer_idx] == 'shared'
-
-        if self.skip_topk:
-            self.indexer = None
-
-    def _get_float_mask(self, query, key, attention_mask, x, attn_mask_type):
-        """Build a FP32 mask with -inf for masked positions."""
-        sq = query.size(0)
-        skv = key.size(0)
-        if attn_mask_type is not None:
-            from megatron.core.transformer.enums import AttnMaskType
-            assert attn_mask_type == AttnMaskType.causal
-            float_mask = torch.triu(
-                torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device),
-                diagonal=1,
-            )
-        else:
-            b = query.size(1)
-            assert attention_mask.shape == (b, 1, sq, skv)
-            mask = attention_mask.squeeze()
-            float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill(mask, float('-inf'))
-        return float_mask
-
-    def forward(
-        self,
-        query,
-        key,
-        value,
-        attention_mask,
-        x,
-        qr,
-        attn_mask_type=None,
-        attention_bias=None,
-        packed_seq_params=None,
-    ):
-        shared_topk_indices = getattr(self, '_shared_topk_indices', None)
-
-        if self.skip_topk:
-            # Shared layer: reuse topk_indices from previous full layer
-            assert shared_topk_indices is not None and 'topk_indices' in shared_topk_indices, (
-                f'GLM 5.2 DSA: Layer {self.layer_number} is a "shared" indexer layer but no '
-                f'"full" layer precedes it in this PP stage. Please adjust '
-                f'`--pipeline_model_parallel_layout` to ensure each PP stage starts with a "full" indexer layer. '
-                f'indexer_types: {self.config.hf_config.indexer_types}.')
-            topk_indices = shared_topk_indices['topk_indices']
-            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
-            return output
-
-        # Full layer: compute topk_indices, store for shared layers, then run sparse attention.
-        # We override the full forward to avoid double-computing topk_indices.
-        x = x.detach()
-        qr = qr.detach()
-        float_mask = self._get_float_mask(query, key, attention_mask, x, attn_mask_type)
-
-        if self.training and torch.is_grad_enabled():
-            q, k, weights = self.indexer.forward_before_topk(x, qr, packed_seq_params)
-            indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0)
-            kwargs = {}
-            if mcore_019:
-                kwargs['calculate_per_token_loss'] = self.config.calculate_per_token_loss
-            topk_indices, indexer_loss = FusedDSAIndexerLoss.apply(
-                q, weights, k, query.detach(), key.detach(), self.softmax_scale,
-                self.indexer.index_topk, indexer_loss_coeff, float_mask,
-                getattr(self.config, 'dsa_indexer_use_sparse_loss', False), self.indexer.pg_collection, **kwargs)
-            if indexer_loss_coeff > 0:
-                DSAIndexerLossLoggingHelper.save_loss_to_tracker(
-                    loss=indexer_loss,
-                    layer_number=self.layer_number,
-                    num_layers=max(
-                        self.layer_number,
-                        self.config.num_layers + (self.config.mtp_num_layers or 0),
-                    ),
-                )
-            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
-            output = DSAIndexerLossAutoScaler.apply(output, indexer_loss)
-        else:
-            _, topk_indices = self.indexer.forward_with_scores(
-                x, qr, mask=float_mask, packed_seq_params=packed_seq_params)
-            output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale)
-
-        # Store topk_indices for subsequent shared layers (in-place dict mutation)
-        if shared_topk_indices is not None:
-            shared_topk_indices['topk_indices'] = topk_indices.detach()
-
-        return output
-
-
-class GlmMoeDsaGPTModel(GPTModel):
-    """GPT model for GLM 5.2 with shared DSA indexer support.
-
-    Creates a ``shared_topk_indices`` dict and passes it through
-    ``extra_block_kwargs`` so that "full" DSA layers can store their
-    topk_indices for reuse by subsequent "shared" layers.
-    """
-
-    def forward(self, *args, **kwargs):
-        extra_block_kwargs = kwargs.get('extra_block_kwargs') or {}
-        extra_block_kwargs['shared_topk_indices'] = {}
-        kwargs['extra_block_kwargs'] = extra_block_kwargs
-        return super().forward(*args, **kwargs)
-
-
-@contextmanager
-def _shared_topk_indices_context(layer, shared_topk_indices):
-    """Temporarily inject shared_topk_indices into the core attention module."""
-    core_attn = None
-    if shared_topk_indices is not None and hasattr(layer, 'self_attention'):
-        _attn = getattr(layer.self_attention, 'core_attention', None)
-        if isinstance(_attn, GlmMoeDsaDSAttention):
-            core_attn = _attn
-            core_attn._shared_topk_indices = shared_topk_indices
-    try:
-        yield
-    finally:
-        if core_attn is not None:
-            core_attn._shared_topk_indices = None
-
-
-class GlmMoeDsaTransformerBlock(TransformerBlock):
-    """TransformerBlock that routes ``shared_topk_indices`` to DSAttention.
-
-    Pops ``shared_topk_indices`` from kwargs before calling the layer
-    (so it doesn't hit ``_forward_attention``'s fixed signature), injects
-    it via context manager, and restores it afterward for subsequent layers.
-    """
-
-    def _layer_forward(self, layer, hidden_states, **kwargs):
-        shared_topk_indices = kwargs.pop('shared_topk_indices', None)
-        with _shared_topk_indices_context(layer, shared_topk_indices):
-            result = super()._layer_forward(layer, hidden_states, **kwargs)
-        # Restore for subsequent layers
-        if shared_topk_indices is not None:
-            kwargs['shared_topk_indices'] = shared_topk_indices
-        return result
-
-
 class GlmMoeDsaLoader(ModelLoader):
-    model_cls = GPTModel if mcore_019 else GlmMoeDsaGPTModel
-    transformer_block = TransformerBlock if mcore_019 else GlmMoeDsaTransformerBlock
 
     def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
+        if not mcore_019:
+            raise ImportError('Please install the megatron-core main branch to support `glm_moe_dsa`:'
+                              '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
         transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
 
         indexer_types = self.config.hf_config.indexer_types
         if indexer_types is not None:
             for layer_spec in transformer_layer_spec.layer_specs:
                 core_attn = layer_spec.submodules.self_attention.submodules.core_attention
-                if hasattr(core_attn, 'module') and issubclass(core_attn.module, DSAttention):
-                    core_attn.module = DSAttention if mcore_019 else GlmMoeDsaDSAttention
+                if hasattr(core_attn, 'module') and issubclass(core_attn.module, McoreDSAttention):
+                    core_attn.module = DSAttention
 
         return transformer_layer_spec
 
diff --git a/src/mcore_bridge/model/modules/absorbed_mla.py b/src/mcore_bridge/model/modules/absorbed_mla.py
index 0a87a41..dadffb4 100644
--- a/src/mcore_bridge/model/modules/absorbed_mla.py
+++ b/src/mcore_bridge/model/modules/absorbed_mla.py
@@ -275,6 +275,16 @@ def forward(
                 packed_seq_params = PackedSeqParams()
             # for easy injection of rotary_pos_emb (patch)
             packed_seq_params.rotary_pos_emb = rotary_pos_emb
+            if self.config.llm_model_type == 'glm_moe_dsa':
+                topk_holder = (
+                    self.core_attention._get_index_share_topk_holder(packed_seq_params, attention_mask)
+                    if self.core_attention.index_share else None)
+                if self.core_attention.skip_topk and self.core_attention.source_layer not in topk_holder:
+                    raise ValueError(
+                        f'DSA: Layer {self.layer_number} is a "shared" indexer layer but no '
+                        f'"full" layer precedes it in this PP stage. Please adjust '
+                        f'`--pipeline_model_parallel_layout` to ensure each PP stage starts with a "full" indexer layer. '
+                        f'indexer_types: {self.config.hf_config.indexer_types}.')
         if self.checkpoint_core_attention and self.training:
             core_attn_out = self._checkpointed_attention_forward(
                 q_absorbed,

From bf23bcb974095149d6ba5500b56ca534209043a7 Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 14 Jul 2026 15:18:58 +0800
Subject: [PATCH 4/7] fix

---
 src/mcore_bridge/model/modules/absorbed_mla.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mcore_bridge/model/modules/absorbed_mla.py b/src/mcore_bridge/model/modules/absorbed_mla.py
index dadffb4..87a5586 100644
--- a/src/mcore_bridge/model/modules/absorbed_mla.py
+++ b/src/mcore_bridge/model/modules/absorbed_mla.py
@@ -2,7 +2,7 @@
     from megatron.core.transformer.experimental_attention_variant.absorbed_mla import \
         AbsorbedMLASelfAttention as McoreAbsorbedMLASelfAttention
 except ImportError:
-    McoreAbsorbedMLASelfAttention = None
+    McoreAbsorbedMLASelfAttention = object
 
 import torch
 from megatron.core import tensor_parallel

From a053583ec4229f3d63c2b39200582eb6e8835bb7 Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 14 Jul 2026 15:43:39 +0800
Subject: [PATCH 5/7] lint pass

---
 src/mcore_bridge/model/modules/absorbed_mla.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/mcore_bridge/model/modules/absorbed_mla.py b/src/mcore_bridge/model/modules/absorbed_mla.py
index 87a5586..c634f9f 100644
--- a/src/mcore_bridge/model/modules/absorbed_mla.py
+++ b/src/mcore_bridge/model/modules/absorbed_mla.py
@@ -280,11 +280,10 @@ def forward(
                     self.core_attention._get_index_share_topk_holder(packed_seq_params, attention_mask)
                     if self.core_attention.index_share else None)
                 if self.core_attention.skip_topk and self.core_attention.source_layer not in topk_holder:
-                    raise ValueError(
-                        f'DSA: Layer {self.layer_number} is a "shared" indexer layer but no '
-                        f'"full" layer precedes it in this PP stage. Please adjust '
-                        f'`--pipeline_model_parallel_layout` to ensure each PP stage starts with a "full" indexer layer. '
-                        f'indexer_types: {self.config.hf_config.indexer_types}.')
+                    raise ValueError(f'DSA: Layer {self.layer_number} is a "shared" indexer layer but no '
+                                     f'"full" layer precedes it in this PP stage. Please adjust '
+                                     f'`--pipeline_model_parallel_layout` to ensure each PP stage starts with '
+                                     f'a "full" indexer layer. indexer_types: {self.config.hf_config.indexer_types}.')
         if self.checkpoint_core_attention and self.training:
             core_attn_out = self._checkpointed_attention_forward(
                 q_absorbed,

From 5a24650a42adae3076ad12c0a4229cf5755ad0de Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 14 Jul 2026 17:24:46 +0800
Subject: [PATCH 6/7] fix

---
 src/mcore_bridge/model/gpts/glm_moe_dsa.py               | 8 ++++----
 src/mcore_bridge/model/modules/multi_latent_attention.py | 4 ++++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
index b98b1a0..c5c1ba7 100644
--- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py
+++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
@@ -29,12 +29,12 @@ def _get_index_share_carrier(self, packed_seq_params, attention_mask):
 class GlmMoeDsaLoader(ModelLoader):
 
     def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
-        if not mcore_019:
-            raise ImportError('Please install the megatron-core main branch to support `glm_moe_dsa`:'
-                              '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
         transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
-
         indexer_types = self.config.hf_config.indexer_types
+        if indexer_types is not None and getattr(DSAttention, '_HOLDER_ATTR', None) is None:
+            raise ImportError(
+                'Please install the megatron-core main branch to support the "shared" indexer layer of `glm_moe_dsa`: '
+                '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
         if indexer_types is not None:
             for layer_spec in transformer_layer_spec.layer_specs:
                 core_attn = layer_spec.submodules.self_attention.submodules.core_attention
diff --git a/src/mcore_bridge/model/modules/multi_latent_attention.py b/src/mcore_bridge/model/modules/multi_latent_attention.py
index 1261157..890f425 100644
--- a/src/mcore_bridge/model/modules/multi_latent_attention.py
+++ b/src/mcore_bridge/model/modules/multi_latent_attention.py
@@ -253,6 +253,10 @@ def forward(
         else:
             extra_kwargs = {}
             if self.config.experimental_attention_variant == 'dsa':
+                if packed_seq_params is not None or self.config.context_parallel_size > 1:
+                    raise ImportError('Please install the megatron-core main branch to support `DSAttention` '
+                                      'padding_free/context parallelism: '
+                                      '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
                 # For dsa we need to pass in the original hidden states and the compressed
                 # query representation.
                 extra_kwargs['x'] = hidden_states

From d4fee1dd3b4502f7cab24412f276e60ab75d93aa Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 14 Jul 2026 17:36:20 +0800
Subject: [PATCH 7/7] fix

---
 src/mcore_bridge/model/gpts/glm_moe_dsa.py | 5 ++---
 src/mcore_bridge/model/register.py         | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/mcore_bridge/model/gpts/glm_moe_dsa.py b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
index c5c1ba7..522be81 100644
--- a/src/mcore_bridge/model/gpts/glm_moe_dsa.py
+++ b/src/mcore_bridge/model/gpts/glm_moe_dsa.py
@@ -30,12 +30,11 @@ class GlmMoeDsaLoader(ModelLoader):
 
     def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
         transformer_layer_spec = super().get_transformer_layer_spec(vp_stage)
-        indexer_types = self.config.hf_config.indexer_types
-        if indexer_types is not None and getattr(DSAttention, '_HOLDER_ATTR', None) is None:
+        if self.config.dsa_indexer_topk_freq > 1 and getattr(DSAttention, '_HOLDER_ATTR', None) is None:
             raise ImportError(
                 'Please install the megatron-core main branch to support the "shared" indexer layer of `glm_moe_dsa`: '
                 '`pip install git+https://github.com/NVIDIA/Megatron-LM.git`')
-        if indexer_types is not None:
+        if self.config.dsa_indexer_topk_freq > 1:
             for layer_spec in transformer_layer_spec.layer_specs:
                 core_attn = layer_spec.submodules.self_attention.submodules.core_attention
                 if hasattr(core_attn, 'module') and issubclass(core_attn.module, McoreDSAttention):
diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py
index 4bf1a7a..e0d6541 100644
--- a/src/mcore_bridge/model/register.py
+++ b/src/mcore_bridge/model/register.py
@@ -163,7 +163,7 @@ def _replace_mla_attention(self, transformer_layer_spec):
             self_attention = layer_spec.submodules.self_attention
             if self_attention.module is McoreMLASelfAttention:
                 self_attention.module = MLASelfAttention
-            elif self_attention.module.__name__ == 'AbsorbedMLASelfAttention':
+            elif getattr(self_attention.module, '__name__', None) == 'AbsorbedMLASelfAttention':
                 self_attention.module = AbsorbedMLASelfAttention
 
     def _replace_router(self, transformer_layer_spec, mlp_key='mlp'):