diff --git a/mindformers/parallel_core/model_parallel_config.py b/mindformers/parallel_core/model_parallel_config.py
index 141728b69f250eb4ff105451b095076ca8fd74c9..fd13659dc721a72b9632e0bdb947a11cb3136ecd 100644
--- a/mindformers/parallel_core/model_parallel_config.py
+++ b/mindformers/parallel_core/model_parallel_config.py
@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from mindformers.tools.utils import get_rank_info
 from mindformers.parallel_core.mf_model_config import convert_str_to_mstype
 
 
@@ -143,3 +144,12 @@ class ModelParallelConfig:
 
         if self.sequence_parallel and self.tensor_model_parallel_size <= 1:
             raise ValueError("Can not use sequence parallelism without tensor parallelism")
+
+    def get_offset_id(self):
+        """get sequence shard offset id"""
+        rank, group_size = get_rank_info()
+        pp_size = group_size // self.pipeline_model_parallel_size
+        dp_size = pp_size // self.data_parallel_size
+        dp_id = rank % pp_size // dp_size
+        cp_id = rank % pp_size % dp_size // self.tensor_model_parallel_size % self.context_parallel_size
+        return dp_id, cp_id
diff --git a/mindformers/parallel_core/training_graph/transformer/dsa/__init__.py b/mindformers/parallel_core/training_graph/transformer/dsa/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/mindformers/parallel_core/training_graph/transformer/dsa/dsa_attention.py b/mindformers/parallel_core/training_graph/transformer/dsa/dsa_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..8642f9a288c82a41c7197b5ac6dc90f41f9d4098
--- /dev/null
+++ b/mindformers/parallel_core/training_graph/transformer/dsa/dsa_attention.py
@@ -0,0 +1,178 @@
+import copy
+from dataclasses import dataclass
+from typing import Union
+from  functools import partial
+
+import mindspore.common.dtype as mstype
+from mindspore import nn, ops, Tensor
+from mindspore.ops import operations as P
+from mindspore.ops import auto_generate as aclnn_ops
+from mindspore.context import ParallelMode
+from mindspore.parallel._utils import _get_parallel_mode
+
+from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module
+from mindformers.parallel_core.transformer_config import TransformerConfig
+from mindformers.parallel_core.training_graph.base_models.common.embeddings.rope_utils import ApplyRotaryPosEmb
+from mindformers.parallel_core.training_graph.transformer.dsa.dsa_indexer_loss import DSAIndexerLoss
+from mindformers.parallel_core.training_graph.device_matrix import layout
+
+
+def sfa_infer_shape(is_tnd, *args):
+    if is_tnd:
+        tq, nq, d = args[0]
+        _, nk, _ = args[1]
+        softmax_shape = (nk, tq, nq // nk)
+    else:
+        b, sq, nq, d = args[0]
+        _, _, nk, _ = args[1]
+        softmax_shape = (b, nk, sq, nq // nk)
+    return args[0], softmax_shape, softmax_shape
+
+
+def sfa_infer_dtype(*args):
+    return args[0], mstype.float32, mstype.float32
+
+
+@dataclass
+class DSAttentionSubmodules:
+    """
+    Configuration class for specifying the submodules of DSAttention.
+
+    Args:
+        indexer: DSA Indexer module for computing sparse attention indices.
+    """
+
+    indexer: Union[ModuleSpec, type] = None
+
+
+class DSAttention(nn.Cell):
+    """
+    This module implements sparse attention mechanism using an DSA Indexer to compute top-k attention indices
+    for reducing computational complexity.
+
+    Reference:
+        https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L491-L597
+    """
+
+    def __init__(
+        self,
+        config: TransformerConfig,
+        submodules: DSAttentionSubmodules,
+        layer_number: int,
+        attention_type: str = None,
+        attn_mask_type: str = None,
+        cp_comm_type: str = None,
+    ):
+        super().__init__(config=config)
+        if attn_mask_type:
+            raise NotImplementedError("For FlashAttention, 'attn_mask_type' is not supported for now.")
+        if attention_type:
+            raise NotImplementedError("For FlashAttention, 'attention_type' is unused for now.")
+        if cp_comm_type:
+            raise NotImplementedError("For FlashAttention, 'cp_comm_type' is not supported for now.")
+
+        self.use_mla_absorb = self.config.use_mla_absorb
+        self.input_layout = self.config.input_layout
+        self.is_tnd = self.input_layout == "TND"
+        self.layer_number = layer_number
+        self.attention_mode = 2 if self.use_mla_absorb else 1
+
+        dp_id, cp_id = config.get_offset_id()
+        self.dp_id = dp_id
+        self.offset_id = cp_id * config.context_parallel_size + dp_id if self.is_tnd else cp_id
+
+        self.indexer = build_module(
+            submodules.indexer, config=self.config
+        )
+        self.softmax_scale = config.softmax_scale or config.kv_channels ** -0.5
+        self.sparse_flash_attention = P.Morph(
+            self._sparse_flash_attention_forward,
+            partial(sfa_infer_shape, self.is_tnd),
+            sfa_infer_dtype).add_prim_attr("self_define_shard", True)
+        self.indexer_loss = DSAIndexerLoss(config, self.softmax_scale)
+
+        self.cast = aclnn_ops.Cast()
+        self.clamp= aclnn_ops.ClampScalar()
+        self.transpose_q = aclnn_ops.Transpose()
+        self.transpose_kv = aclnn_ops.Transpose()
+        self.transpose_attn = aclnn_ops.Transpose()
+        self.split_q = aclnn_ops.SplitWithSize()
+        self.split_k = aclnn_ops.SplitWithSize()
+
+        if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,):
+            self.shard()
+
+    def _sparse_flash_attention_forward(
+        self, q, k, v, topk_indices,
+        q_rope=None, k_rope=None,
+        actual_seq_qlen=None, actual_seq_kvlen=None
+    ):
+        """get attention output with fused sfa operator, only support attention mode = 2."""
+        if self.is_tnd:
+            slice_tq = q.shape[0]
+            slice_tk = k.shape[0]
+            offset_q = slice_tq * self.offset_id
+            offset_k = slice_tk * self.dp_id
+            actual_seq_qlen = self.cast(self.clamp(actual_seq_qlen - offset_q, 0, slice_tq), mstype.int32)
+            actual_seq_kvlen = self.cast(self.clamp(actual_seq_kvlen - offset_k, 0, slice_tk), mstype.int32)
+
+        attention_out, softmax_max, softmax_sum = ops.sparse_flash_attention(
+            q, k, v, topk_indices, self.softmax_scale,
+            q_rope=q_rope, k_rope=k_rope,
+            actual_seq_lengths_query=actual_seq_qlen,
+            actual_seq_lengths_kv=actual_seq_kvlen,
+            layout_query=self.input_layout,
+            layout_kv=self.input_layout,
+            attention_mode=self.attention_mode,
+            return_softmax_lse=True
+        )
+        return attention_out, softmax_max, softmax_sum
+
+    def construct(
+            self, query, key, value, topk_indices,
+            actual_seq_qlen=None, actual_seq_kvlen=None
+    ):
+        """Forward pass for Sparse Attention."""
+        if not self.is_tnd:
+            query = self.transpose_q(query, (1, 0, 2, 3))
+            key = self.transpose_kv(key, (1, 0, 2, 3))
+            value = self.transpose_kv(value, (1, 0, 2, 3))
+
+        q_nope, q_rope = self.split_q(query, [self.kv_lora_rank, self.qk_pos_emb_head_dim], -1)
+        k_nope, k_rope = self.split_k(key, [self.kv_lora_rank, self.qk_pos_emb_head_dim], -1)
+
+        # ===================================
+        # Run sparse attention kernel
+        # ===================================
+        attention_output = self.sparse_flash_attention(
+            q_nope, k_nope, value, topk_indices,
+            q_rope, k_rope, actual_seq_qlen, actual_seq_kvlen
+        )[0]
+        if not self.is_tnd:
+            attention_output = self.transpose_attn(attention_output, (1, 0, 2, 3))
+
+        return attention_output
+
+    def shard(self):
+        kv_tp_shard = "None" if self.use_mla_absorb else "tp"
+        self.transpose_q.shard((layout("cp", "dp", "tp", "None"),))
+        self.transpose_kv.shard((layout("cp", "dp", kv_tp_shard, "None"),))
+        self.transpose_attn.shard((layout("dp", "cp", "tp", "None"),))
+        if not self.is_tnd:
+            q_shard = layout("dp", "cp", "tp", "None")
+            kv_shard = layout("dp", "None", kv_tp_shard, "None")
+            idx_shard = layout("dp", "cp", "None", "None")
+            sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard)
+            attn_shard = layout("dp", "cp", "tp", "None")
+            softmax_shard = layout("dp", kv_tp_shard, "cp", "None")
+        else:
+            q_shard = layout("dp_cp", "tp", "None")
+            kv_shard = layout("dp", kv_tp_shard, "None")
+            idx_shard = layout("dp_cp", "None", "None")
+            sfa_shard = (q_shard, kv_shard, kv_shard, idx_shard, q_shard, kv_shard, layout("dp"), layout("dp"))
+            attn_shard = layout("dp_cp", "tp", "None")
+            softmax_shard = layout(kv_tp_shard, "cp", "None")
+        self.split_q.shard((q_shard,))
+        self.split_kv.shard((kv_shard,))
+        self.sparse_flash_attention.shard(sfa_shard, (attn_shard, softmax_shard, softmax_shard))
+
diff --git a/mindformers/parallel_core/training_graph/transformer/dsa/dsa_indexer.py b/mindformers/parallel_core/training_graph/transformer/dsa/dsa_indexer.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d6161cb5a78d65e2972cc4df3ecc324f9350dbf
--- /dev/null
+++ b/mindformers/parallel_core/training_graph/transformer/dsa/dsa_indexer.py
@@ -0,0 +1,362 @@
+import copy
+from dataclasses import dataclass
+from typing import Union
+from  functools import partial
+
+import mindspore.common.dtype as mstype
+from mindspore import nn, ops, Tensor
+from mindspore.ops import operations as P
+from mindspore.ops import auto_generate as aclnn_ops
+from mindspore.context import ParallelMode
+from mindspore.parallel._utils import _get_parallel_mode
+
+from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module
+from mindformers.parallel_core.transformer_config import MLATransformerConfig
+from mindformers.parallel_core.training_graph.base_models.common.embeddings.rope_utils import ApplyRotaryPosEmb
+from mindformers.parallel_core.training_graph.device_matrix import layout
+from mindformers.parallel_core.training_graph.transformer.dsa.hadamard import Hadamard
+
+
+def indexer_infer_shape(sparse_count, *args):
+    q_shape = args[0]
+    k_shape = args[1]
+    infer_shape = q_shape[:-2] + (k_shape[1], sparse_count)
+    return infer_shape, infer_shape
+
+
+def indexer_infer_dtype(*args):
+    return mstype.int32, args[0]
+
+
+@dataclass
+class DSAIndexerSubmodules:
+    """
+    Configuration class for specifying the submodules of an DSA Indexer.
+
+    Args:
+        linear_wq_b: Linear projection for query bottleneck expansion.
+        linear_wk: Linear projection for key.
+        k_norm: Layer normalization for key.
+        linear_weights_proj: Linear projection for attention weights.
+    """
+
+    linear_wq_b: Union[ModuleSpec, type] = None
+    linear_wk: Union[ModuleSpec, type] = None
+    k_norm: Union[ModuleSpec, type] = None
+    linear_weights_proj: Union[ModuleSpec, type] = None
+
+class DSAIndexer(nn.Cell):
+    """
+    DSA Lightning Indexer for DeepSeek Sparse Attention.
+    Computes index scores to identify the top-k most relevant key-value pairs for each query in sparse attention.
+
+    Reference:
+        https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L431-L480
+
+    Args:
+        config (TransformerConfig): The configuration for the transformer model.
+        submodules (DSAIndexerSubmodules): Indexer submodules specification.
+    """
+
+    def __init__(
+        self,
+        config: MLATransformerConfig,
+        submodules: DSAIndexerSubmodules,
+    ) -> None:
+        """Initialize the indexer."""
+        super().__init__()
+        self.config = config
+        self.cp = self.config.context_parallel_size
+        self.input_layout = self.config.input_layout
+        self.is_tnd = self.input_layout == "TND"
+        self.hidden_size = self.config.hidden_size
+        self.qk_pos_emb_head_dim = self.config.qk_pos_emb_head_dim
+        self.q_lora_rank = (
+            self.config.q_lora_rank
+            if self.config.q_lora_rank is not None
+            else self.config.hidden_size
+        )
+
+        dp_id, cp_id = config.get_offset_id()
+        self.dp_id = dp_id
+        cp_id = 0
+        self.offset_id = cp_id * config.context_parallel_size + dp_id if self.is_tnd else cp_id
+
+        self.index_n_heads = self.config.dsa_indexer_n_heads
+        self.index_head_dim = self.config.dsa_indexer_head_dim
+        self.index_topk = self.config.dsa_indexer_topk
+        self.use_sparse_loss = self.config.dsa_indexer_use_sparse_loss
+        self.softmax_scale: float = self.index_head_dim ** -0.5
+
+        self.apply_rotary_emb_q = ApplyRotaryPosEmb(config)
+        self.apply_rotary_emb_k = ApplyRotaryPosEmb(config, for_k_pos_emb=True)
+        self.linear_wq_b = build_module(
+            submodules.linear_wq_b,
+            input_size=self.q_lora_rank,
+            output_size=self.index_n_heads * self.index_head_dim,
+            config=self.config,
+            init_method=self.config.init_method,
+            bias=False,
+            skip_bias_add=False,
+            skip_weight_param_allocation=False,
+        )
+
+        self.linear_wk = build_module(
+            submodules.linear_wk,
+            self.hidden_size,
+            self.index_head_dim,
+            config=self.config,
+            init_method=self.config.init_method,
+            bias=False,
+            skip_bias_add=False,
+            skip_weight_param_allocation=False,
+        )
+
+        k_norm_config = copy.copy(self.config)
+        k_norm_config.normalization = "LayerNorm"
+        self.k_norm = build_module(
+            submodules.k_norm,
+            config=k_norm_config,
+            hidden_size=self.index_head_dim,
+            eps=self.config.layernorm_epsilon,
+        )
+
+        self.linear_weights_proj = build_module(
+            submodules.linear_weights_proj,
+            self.hidden_size,
+            self.index_n_heads,
+            config=self.config,
+            init_method=self.config.init_method,
+            bias=False,
+            skip_bias_add=False,
+            skip_weight_param_allocation=False,
+        )
+
+        self.split_q = aclnn_ops.SplitWithSize()
+        self.split_k = aclnn_ops.SplitWithSize()
+        self.cat_q = aclnn_ops.Concat(-1)
+        self.cat_k = aclnn_ops.Concat(-1)
+        self.cast = aclnn_ops.Cast()
+        self.clamp= aclnn_ops.ClampScalar()
+        self.einsum = aclnn_ops.EinsumExt('bshd,btd->bsht')
+        self.relu = aclnn_ops.ReLU()
+        self.mul = aclnn_ops.Mul()
+        self.sum = aclnn_ops.SumExt()
+        self.reshape = aclnn_ops.Reshape()
+        self.expand_dims = aclnn_ops.ExpandDims()
+        self.transpose_q = aclnn_ops.Transpose()
+        self.transpose_k = aclnn_ops.Transpose()
+        self.transpose_w = aclnn_ops.Transpose()
+        self.topk = aclnn_ops.TopK()
+        self.slice = aclnn_ops.SliceExt()
+        self.compute_sparse_indices = P.Morph(self._compute_sparse_indices_forward,
+                                              partial(indexer_infer_shape, self.index_topk),
+                                              indexer_infer_dtype).add_prim_attr("self_define_shard", True)
+        self.hadamard_q = Hadamard(config)
+        self.hadamard_k = Hadamard(config, for_k_emb=True)
+
+        self.inf = Tensor([-10000], mstype.float32)
+        if _get_parallel_mode() in (ParallelMode.SEMI_AUTO_PARALLEL,):
+            self.shard()
+
+    def convert_layout(self, q, k, weights):
+        """
+        Convert q, k, weights from SB* layout to BS* or T* layout.
+        """
+        seqlen, bsz, num_head, dim_head = q.shape
+        q = self.transpose_q(q, (1, 0, 2, 3))
+        k = self.transpose_k(k, (1, 0, 2, 3))
+        weights = self.transpose_w(weights, (1, 0, 2))
+        if self.is_tnd:
+            q = self.reshape(q, (seqlen * bsz, num_head, dim_head))
+            k = self.reshape(k, (seqlen * bsz, 1, dim_head))
+            weights = self.reshape(weights, (seqlen * bsz, num_head, 1))
+        else:
+            if self.use_sparse_loss:
+                k = self.reshape(k, (bsz, seqlen, 1, dim_head))
+            else:
+                k = self.reshape(k, (bsz, seqlen, dim_head))
+            weights = self.reshape(weights, (seqlen, bsz, num_head, 1))
+        return q, k, weights
+
+    def _compute_tnd_dense_indices(self, q, k, weights, mask, actual_seq_qlen, actual_seq_klen):
+        """Compute full index score and topk indices of TND layout with unfused operators. NOT IMPLEMENTED YET!"""
+        raise NotImplemented
+
+    def _compute_bsnd_dense_indices(self, q, k, weights, mask):
+        """Compute full index score and topk indices of BSND layout with unfused operators."""
+        bsz, seqlen_q = q.shape[:2]
+        seqlen_k = k.shape[1]
+        # Compute attention scores: q @ k^T
+        # [batch, seqlen_q, index_n_heads, index_head_dim] @ [batch, seqlen_k, index_head_dim]^T
+        #   -> [batch, seqlen_q, index_n_heads, seqlen_k]
+        index_scores = self.einsum(q, k)
+
+        # Apply ReLU activation.
+        index_scores = self.relu(index_scores)
+
+        # Weight each head by attention weights.
+        # [batch, seqlen_q, index_n_heads, seqlen_k] * [batch, seqlen_q, index_n_heads, 1]
+        #   -> [batch, seqlen_q, index_n_heads, seqlen_k]
+        index_scores = self.mul(index_scores, weights)
+
+        # Sum across attention heads.
+        # [batch, seqlen_q, index_n_heads, seqlen_k] -> [batch, seqlen_q, 1, seqlen_k]
+        mask = self.reshape(mask, (bsz, seqlen_q, 1, -1))
+        index_scores = self.sum(index_scores, 2, True)
+        index_scores = self.add(index_scores, mask * self.inf)
+
+        topk_indices = self.topk(index_scores, self.index_topk, -1)[1]
+        mask_sliced = self.cast(self.slice(mask, -1, 0, self.sparse_count, 1) * -seqlen_k, mstype.int32)
+        topk_indices = self.clamp(self.add(topk_indices, mask_sliced), -1)
+        return topk_indices, index_scores
+
+    def _compute_dense_indices(self, q, k, weights, mask, actual_seq_qlen=None, actual_seq_klen=None):
+        """Compute full index score and topk indices with unfused operators in dense training stage."""
+        if self.is_tnd:
+            return self._compute_tnd_dense_indices(q, k, weights, mask, actual_seq_qlen, actual_seq_klen)
+        return self._compute_bsnd_dense_indices(q, k, weights, mask)
+
+    def _compute_sparse_indices_forward(self, q, k, weights, actual_seq_qlen=None, actual_seq_klen=None):
+        """Compute topk index score and indices with fused operator in sparse training stage."""
+        if self.is_tnd:
+            slice_tq = q.shape[0]
+            slice_tk = k.shape[0]
+            offset_q = slice_tq * self.offset_id
+            offset_k = slice_tk * self.dp_id
+            actual_seq_qlen = self.cast(self.clamp(actual_seq_qlen - offset_q, 0, slice_tq), mstype.int32)
+            actual_seq_klen = self.cast(self.clamp(actual_seq_klen - offset_k, 0, slice_tk), mstype.int32)
+
+        return ops.lightning_indexer(
+            q, k, weights,
+            actual_seq_lengths_query=actual_seq_qlen,
+            actual_seq_lengths_key=actual_seq_klen,
+            layout_query=self.input_layout,
+            layout_key=self.input_layout,
+            sparse_count=self.index_topk,
+            return_value=True
+        )
+
+    def _compute_index_scores(self, q, k, weights, mask, actual_seq_qlen=None, actual_seq_klen=None):
+        """"""
+        if self.use_sparse_loss:
+            return self.compute_sparse_indices(q, k, weights, actual_seq_qlen, actual_seq_klen)
+        else:
+            return self._compute_dense_indices(q, k, weights, mask, actual_seq_qlen, actual_seq_klen)
+
+    def construct(self, x, qr, mask, rotary_pos_emb, actual_seq_qlen=None, actual_seq_klen=None):
+        """
+        Forward pass for DSA Indexer that returns both index scores and top-k indices.
+        This is used when KL loss is enabled to compare indexer scores with true attention scores.
+        """
+        # =========================================
+        # Get sequence length and batch size
+        # =========================================
+        seqlen, bsz, _ = x.size()
+
+        # =========================================
+        # q linear and apply rope to q
+        # =========================================
+        # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim]
+        q, _ = self.linear_wq_b(qr)
+        # [seqlen, batch, index_n_heads * index_head_dim]
+        #   -> [seqlen, batch, index_n_heads, index_head_dim]
+        q = self.reshape(q, (seqlen, bsz, self.index_n_heads, self.index_head_dim))
+        q_nope, q_pe = self.split_q(
+            q, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1
+        )
+        q_pe = self.apply_rotary_emb_q(
+            q_pe,
+            rotary_pos_emb,
+            rotary_interleaved=self.config.rotary_interleaved,
+            multi_latent_attention=self.config.multi_latent_attention
+        )
+        # [seqlen, batch, *, index_head_dim]
+        q = self.cat_q([q_nope, q_pe])
+
+        # =========================================
+        # k linear and apply rope to k
+        # =========================================
+        # [seqlen, batch, hidden_size] -> [seqlen, batch, 1, index_head_dim]
+        k, _ = self.linear_wk(x)
+        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)
+        # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim]
+        k_nope, k_pe = self.split_k(
+            k, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1
+        )
+        k_pe = self.apply_rotary_emb_k(
+            k_pe,
+            rotary_pos_emb,
+            rotary_interleaved=self.config.rotary_interleaved,
+            multi_latent_attention=self.config.multi_latent_attention
+        )
+        # [seqlen, batch, *, index_head_dim]
+        k = self.cat_k([k_nope, k_pe])
+
+        # =========================================
+        # Rotate activation
+        # =========================================
+        q = self.hadamard_q(q)
+        k = self.hadamard_k(k)
+
+        # =========================================
+        # Compute index scores
+        # =========================================
+        # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads]
+        weights, _ = self.linear_weights_proj(x)
+        weights = weights * (self.index_n_heads ** -0.5) * self.softmax_scale
+        # [batch, seqlen, seqlen]
+        q, k, weights = self.convert_layout(q, k, weights)
+        topk_indices, index_scores = self._compute_index_scores(q, k, weights, mask, actual_seq_qlen, actual_seq_klen)
+        return index_scores, topk_indices
+
+    def shard(self):
+        """sharding operators"""
+        # before convert_layout
+        q_shard = layout("cp", "dp", "tp", "None")
+        k_shard = layout("cp", "dp", "None", "None")
+        w_shard = layout("cp", "dp", "tp")
+        self.split_q.shard((q_shard,))
+        self.split_k.shard((k_shard,))
+        self.cat_q.shard((q_shard, q_shard))
+        self.cat_k.shard((k_shard, k_shard))
+        self.transpose_q.shard((q_shard,))
+        self.transpose_k.shard((k_shard,))
+        self.transpose_w.shard((w_shard,))
+
+        # after convert_layout
+        if self.use_sparse_loss:
+            if self.is_tnd:
+                q_shard = layout("dp", "None", "None")
+                k_shard = layout("dp", "None", "None")
+                w_shard = layout("dp", "None")
+                actual_seq_shard = layout("dp")
+                indexer_shard = layout("dp", "None", "None")
+                self.compute_sparse_indices.shard(
+                    (q_shard, k_shard, w_shard, actual_seq_shard, actual_seq_shard), (indexer_shard, indexer_shard)
+                )
+            else:
+                q_shard = layout("dp", "None", "None", "None")
+                k_shard = layout("dp", "None", "None", "None")
+                w_shard = layout("dp", "None", "None")
+                indexer_shard = layout("dp", "None", "None", "None")
+                self.compute_sparse_indices.shard(
+                    (q_shard, k_shard, w_shard), (indexer_shard, indexer_shard)
+                )
+        else:
+            q_shard = layout("dp", "cp", "tp", "None")
+            k_shard = layout("dp", "cp", "None")
+            w_shard = layout("dp", "cp", "tp", "None")
+            qk_shard = layout("dp", "cp", "tp", "None")
+            sum_shard = layout("dp", "cp", "None", "None")
+            score_shard = layout("dp", "cp", "None", "None")
+            self.einsum.shard((q_shard, k_shard), (qk_shard,))
+            self.relu.shard((qk_shard,))
+            self.mul.shard((qk_shard, w_shard))
+            self.sum.shard((sum_shard,))
+            self.add.shard((score_shard, score_shard))
+            self.topk.shard((score_shard,))
+            self.slice.shard((score_shard,))
+            self.clamp.shard((score_shard,))
\ No newline at end of file
diff --git a/mindformers/parallel_core/training_graph/transformer/dsa/dsa_indexer_loss.py b/mindformers/parallel_core/training_graph/transformer/dsa/dsa_indexer_loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..56f2fd27e02ec2d46dd6f4eae81446dd540e04d6
--- /dev/null
+++ b/mindformers/parallel_core/training_graph/transformer/dsa/dsa_indexer_loss.py
@@ -0,0 +1,126 @@
+import mindspore.common.dtype as mstype
+from mindspore import nn, Tensor
+from mindspore.ops import auto_generate as aclnn_ops
+from mindspore.context import ParallelMode
+from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation
+
+from mindformers.parallel_core.transformer_config import TransformerConfig
+from mindformers.parallel_core.training_graph.device_matrix import layout
+
+
+class DSAIndexerLoss(nn.Cell):
+    """
+    Compute KL divergence loss between index_scores and true attention_scores.
+
+    This loss trains the indexer to predict which tokens are important by matching the distribution
+    of true attention scores.
+
+    Reference:
+        Section 2.1 of https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf
+
+    Args:
+        config (TransformerConfig): The configuration for the transformer model.
+        softmax_scale: Scale coefficient after q @ k^T.
+
+    Input:
+        index_scores: Scores predicted by indexer [batch, seqlen_q, seqlen_k].
+        topk_indices: Top-k indices [batch, seqlen_q, index_topk].
+        query: Query tensor [seqlen_q, batch, heads, dim].
+        key: Key tensor [seqlen_k, batch, heads, dim].
+        softmax_scale: Scale coefficient after q @ k^T.
+        loss_coeff: Coefficient for the indexer KL divergence loss.
+        sparse_loss: bool, whether to use sparse indexer loss. If True, only the topk
+            indices will be used to compute the loss.
+        pg_collection: Process group collection, must have TP process group.
+
+    Returns:
+        index_loss: KL divergence loss (scalar).
+    """
+    def __init__(self, config: TransformerConfig, softmax_scale=None):
+        super().__init__()
+        self.use_sparse_loss = config.dsa_indexer_use_sparse_loss
+        self.loss_coeff = config.dsa_indexer_loss_coeff
+        self.sparse_count = config.dsa_indexer_topk
+        self.num_attention_heads = config.num_attention_heads
+        self.is_tnd = config.input_layout == "TND"
+
+        self.softmax_scale = softmax_scale or config.kv_channels ** -0.5
+        self.inf = Tensor([-10000], mstype.float32)
+
+        self.transpose = aclnn_ops.Transpose()
+        self.bmm = aclnn_ops.BatchMatMulExt()
+        self.softmax1 = aclnn_ops.Softmax(-1)
+        self.softmax2 = aclnn_ops.Softmax(-1)
+        self.cast = aclnn_ops.Cast()
+        self.reshape = aclnn_ops.Reshape()
+        self.gather = aclnn_ops.Gather(3)
+        self.tile = aclnn_ops.Tile()
+        self.slice = aclnn_ops.SliceExt()
+        self.sum1 = aclnn_ops.SumExt()
+        self.div = aclnn_ops.Div()
+        self.sum2 = aclnn_ops.SumExt()
+        self.kl_div = aclnn_ops.KLDiv()
+        self.sum3 = aclnn_ops.SumExt()
+        self.mean = aclnn_ops.MeanExt()
+
+    def construct(self, index_scores, topk_indices, query, key, mask):
+        bsz, seqlen, num_head, hidden_size = query.shape
+        # q: [sq, b, n, d] ->  [b, n, sq, d], k: [sk, b, n, d] ->  [b, n, d, sk]
+        # attention_score: q @ k -> [b, n, sq, sk]
+        query = self.cast(self.transpose(query, (1, 2, 0, 3)), mstype.float32)
+        key = self.cast(self.transpose(key, (1, 2, 3, 0)), mstype.float32)
+        attention_scores = self.bmm(query, key) * self.softmax_scale
+
+        if self.use_sparse_loss:
+            # when use sparse loss, the input `index_scores` is already in topk format
+            # topk_indices: [b, sq, 1, sk] -> [b, n, sq, sk]
+            # attention_score: [b, n, sq, sk] -> [b, n, sq, topk] -> [b, sq, topk]
+            topk_indices = self.reshape(topk_indices, (bsz, 1, seqlen, -1))
+            topk_indices = self.tile(topk_indices, (1, self.num_attention_heads, 1, 1))
+            attention_scores = self.gather(attention_scores, topk_indices, -1)
+            mask = self.cast(mask, mstype.float32)
+            mask = self.slice(mask, -1, 0, self.sparse_count, 1) * self.inf
+
+        attention_scores = self.sub(attention_scores, mask)
+        attention_scores = self.softmax1(attention_scores)
+        attention_scores = self.sum1(attention_scores, dim=1)
+        attention_scores = self.div(attention_scores, self.sum2(attention_scores, -1, keepdim=True))
+
+        # compute KL divergence loss. index_scores is already masked.
+        index_scores = self.reshape(index_scores, (bsz, seqlen, -1))
+        index_scores = self.cast(index_scores, mstype.float32)
+        index_scores = self.softmax2(index_scores)
+
+        kl_per_element = self.kl_div(index_scores, attention_scores, "none")
+
+        # [b, sq, sk] -> [b, sq] -> [1]
+        # Each token has same weight in the loss.
+        kl_div = self.mean(self.sum3(kl_per_element, -1))
+
+        # Scale by coefficient.
+        indexer_loss = kl_div * self.loss_coeff
+
+        return indexer_loss
+
+    def shard(self):
+        input_shard = layout("cp", "dp", "tp", "None")
+        q_shard = layout("dp", "tp", "cp", "None")
+        k_shard = layout("dp", "tp", "None", "None")
+        attn_shard = q_shard
+        idx_shard = layout("dp", "None", "cp", "None")
+        mask_shard = idx_shard
+        pooled_attn_shard = layout("dp", "cp", "None")
+        self.transpose.shard((input_shard,))
+        self.bmm.shard((q_shard, k_shard))
+        self.tile.shard((idx_shard,), (attn_shard,))
+        self.gather.shard((attn_shard, attn_shard))
+        self.slice.shard((mask_shard,))
+        self.sub.shard((attn_shard, attn_shard))
+        self.softmax1.shard((attn_shard,))
+        self.sum1.shard((layout("dp", "None", "cp", "None"),))
+        self.softmax2.shard((pooled_attn_shard,))
+        self.sum2.shard((pooled_attn_shard,))
+        self.div.shard((pooled_attn_shard, pooled_attn_shard))
+        self.kl_div.shard((pooled_attn_shard, pooled_attn_shard))
+        self.sum3.shard((pooled_attn_shard,))
+        self.mean.shard((layout("dp", "cp"),))
diff --git a/mindformers/parallel_core/training_graph/transformer/dsa/hadamard.py b/mindformers/parallel_core/training_graph/transformer/dsa/hadamard.py
new file mode 100644
index 0000000000000000000000000000000000000000..f92f26b6de6e891ea4bb022449180bafbcd32f99
--- /dev/null
+++ b/mindformers/parallel_core/training_graph/transformer/dsa/hadamard.py
@@ -0,0 +1,31 @@
+from scipy.linalg import hadamard
+
+from mindspore.ops import auto_generate as aclnn_ops
+from mindspore import nn, Tensor
+import mindspore.common.dtype as mstype
+
+from mindformers.parallel_core.transformer_config import TransformerConfig
+from mindformers.parallel_core.training_graph.device_matrix import layout
+
+
+class Hadamard(nn.Cell):
+    def __init__(self, config: TransformerConfig, for_k_emb=False):
+        super().__init__()
+        self.hadamard_mat = Tensor(hadamard(config.dsa_indexer_head_dim), mstype.float32)
+        self.scale = config.dsa_indexer_head_dim ** -0.5
+        self.for_k_emb = for_k_emb
+        self.bmm = aclnn_ops.BatchMatMul()
+        self.cast = aclnn_ops.Cast()
+        self.mul = aclnn_ops.Mul()
+        self.shard()
+
+    def construct(self, x):
+        hadamard_mat = self.cast(self.hadamard_mat, x.dtype)
+        x = self.bmm(x, hadamard_mat)
+        x = self.mul(x, self.scale)
+        return x
+
+    def shard(self):
+        bmm_shard = layout("cp", "dp", "None", "None") if self.for_k_emb else layout("cp", "dp", "tp", "None")
+        self.bmm.shard((bmm_shard, layout("None", "None")))
+        self.mul.shard((bmm_shard, layout("None")))
\ No newline at end of file
diff --git a/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py b/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py
index 18f83d65d72dcbe24f0e0107bc1d68b98a7256ec..e4703e2bd4eb6433dad5a94217ac8f57b8ff214d 100644
--- a/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py
+++ b/mindformers/parallel_core/training_graph/transformer/multi_latent_attention.py
@@ -18,7 +18,17 @@ from typing import Union
 import math
 
 from mindspore import nn, Tensor
-from mindspore.ops.auto_generate import Cast, Shape, Reshape, Transpose, Tile, Concat, SplitWithSize, ExpandDims
+from mindspore.ops.auto_generate import (
+    Cast,
+    Shape,
+    Reshape,
+    Transpose,
+    Tile,
+    Concat,
+    SplitWithSize,
+    ExpandDims,
+    EinsumExt
+)
 from mindspore.ops import functional as F
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_parallel_mode, _is_sharding_propagation
@@ -107,6 +117,8 @@ class MultiLatentAttention(nn.Cell):
         self.v_head_dim = self.config.v_head_dim
         self.input_layout = self.config.input_layout
         self.compute_dtype = self.config.compute_dtype
+        self.use_mla_absorb = self.config.use_mla_absorb
+        self.use_dsa = self.config.experimental_attention_variant == "dsa"
 
         # Define ulysses context parallel related parameters
         self.cp_ds = self.config.hierarchical_context_parallel_sizes
@@ -130,6 +142,9 @@ class MultiLatentAttention(nn.Cell):
             layer_number=self.layer_index,
             softmax_scale=self.softmax_scale,
         )
+        if self.use_dsa:
+            self.indexer = self.core_attention.indexer
+            self.indexer_loss = self.core_attention.indexer_loss
 
         # Output.
         self.linear_proj = build_module(
@@ -148,7 +163,15 @@ class MultiLatentAttention(nn.Cell):
         self.reshape = Reshape()
         self.bs_transpose = Transpose()
         self.tnd_transpose = Transpose()
+        self.tnd_transpose_one_head = Transpose()
         self.cast = Cast()
+        if self.use_mla_absorb:
+            self.split_w_kvb = SplitWithSize()
+            self.transpose_weight = Transpose()
+            self.w_kvb = self.linear_kvb.weight if self.config.mla_qkv_concat else self.linear_kv_up_proj.weight
+            self.cat_q = Concat(axis=3)
+            self.cat_k = Concat(axis=3)
+            self.einsum = EinsumExt("sbnd,ndr->sbnr")
 
         if _get_parallel_mode() in (ParallelMode.AUTO_PARALLEL,) and _is_sharding_propagation():
             self.sharding_propagation()
@@ -183,6 +206,13 @@ class MultiLatentAttention(nn.Cell):
 
         self.bs_transpose.shard(((dp, cp, tp),))
         self.tnd_transpose.shard((layout("cp", "dp", "tp", "None"),))
+        self.tnd_transpose_one_head.shard((layout("cp", "dp", "None", "None"),))
+        if self.use_mla_absorb:
+            self.split_w_kvb.shard((layout("tp", "None", "None"),))
+            self.transpose_weight.shard((layout("tp", "None", "None"),))
+            self.cat_q.shard((layout("cp", "dp", "tp", "None"), layout("cp", "dp", "tp", "None")))
+            self.cat_k.shard((layout("cp", "dp", "None", "None"), layout("cp", "dp", "None", "None")))
+            self.einsum.shard((layout("cp", "dp", "tp", "None"), layout("tp", "None", "None")))
 
     def construct(self, x: Tensor, attention_mask=None, rotary_pos_emb=None, rotary_pos_cos=None,
                   rotary_pos_sin=None, prefix_keys_values=None, pad_zeros=None, actual_seq_len=None):
@@ -195,11 +225,29 @@ class MultiLatentAttention(nn.Cell):
             raise NotImplementedError("prefix_keys_values is not supported for now.")
         ori_dtype = x.dtype
         seq_len, bs, _ = self.shape(x)
-
-        query, key, value = self.get_query_key_value_tensors(x, rotary_pos_emb=rotary_pos_emb)
-
+        if self.use_dsa:
+            query, key, value, q_compress = self.get_query_key_value_tensors(x, rotary_pos_emb=rotary_pos_emb)
+            index_scores, topk_indices = self.indexer(
+                F.stop_gradient(x), F.stop_gradient(q_compress),
+                attention_mask, rotary_pos_emb,
+                actual_seq_len, actual_seq_len
+            )
+        else:
+            query, key, value = self.get_query_key_value_tensors(x, rotary_pos_emb=rotary_pos_emb)
+
+        if self.use_mla_absorb:
+            q_nope, q_rope = query
+            k_nope, k_rope = key
+            w_kvb = self.cast(self.w_kvb, self.compute_dtype)
+            w_kvb = self.reshape(w_kvb, (self.num_attention_heads, -1, self.kv_lora_rank))
+            q_absorb, v_absorb = self.split_w_kvb(w_kvb, [self.qk_head_dim, self.v_head_dim], dim=1)
+            v_absorb = self.transpose_weight(v_absorb, (0, 2, 1))
+            q_nope = self.einsum(q_nope, q_absorb)
+            query = self.cat_q([q_nope, q_rope])
+            key = self.cat_k([k_nope, k_rope])
+        head_dim = query.shape[-1]
         # with ulysses context parallel, insert all to all before FA
-        if self.cp > 1:
+        if self.cp > 1 and not self.use_dsa:
             if self.cp_ds > 1:
                 # For query & key & value, transpose from [S, B, N, D] back to [B, S, N, D]
                 query = self._ulysses_qkv_a2a(query)
@@ -213,13 +261,15 @@ class MultiLatentAttention(nn.Cell):
 
         if self.input_layout == "TND":
             query = self.sbh2tnd(query)
-            key = self.sbh2tnd(key)
-            value = self.sbh2tnd(value)
+            key = self.sbh2tnd(key, self.use_mla_absorb)
+            value = self.sbh2tnd(value, self.use_mla_absorb)
 
         query = self.cast(query, self.compute_dtype)
         key = self.cast(key, self.compute_dtype)
         value = self.cast(value, self.compute_dtype)
         if self.use_flash_attention:
+            if self.use_dsa:
+                attention_mask = topk_indices
             if self.use_eod_attn_mask_compression:
                 context_layer = self.core_attention(
                     query, key, value, attention_mask,
@@ -236,12 +286,27 @@ class MultiLatentAttention(nn.Cell):
                 attn_out = self.reshape(context_layer, (seq_len, bs, -1))
         else:
             attn_out = self.core_attention(query, key, value, attention_mask)
-
+        if self.use_mla_absorb:
+            attn_out = self.reshape(attn_out, (seq_len, bs, self.num_attention_heads, -1))
+            attn_out = self.einsum(attn_out, v_absorb)
+            attn_out = self.reshape(attn_out, (seq_len, bs, -1))
         output = self.linear_proj(attn_out)[0]  # dp, mp -> dp, 1 / dp * mp, 1
         output = self.cast(output, ori_dtype)
+        if self.use_dsa:
+            query = F.stop_gradient(query)
+            key = F.stop_gradient(key)
+            if self.use_eod_attn_mask_compression:
+                query = self.reshape(query, (bs, seq_len, -1))
+                key = self.reshape(key, (bs, seq_len, -1))
+                query = self.bs_transpose(query, (1, 0, 2))
+                key = self.bs_transpose(key, (1, 0, 2))
+            query = self.reshape(query, (seq_len, bs, -1, head_dim))
+            key = self.reshape(key, (seq_len, bs, -1, head_dim))
+            indexer_loss = self.indexer_loss(index_scores, topk_indices, query, key, attention_mask)
+            # return output, indexer_loss
         return output
 
-    def sbh2tnd(self, x):
+    def sbh2tnd(self, x, one_head=False):
         """
         Convert a input tensor from SBH/SBND layout to TND layout.
 
@@ -252,9 +317,10 @@ class MultiLatentAttention(nn.Cell):
             x_merge: the TND output tensor
         """
         seq_len, bs = x.shape[:2]
-        x = self.reshape(x, (seq_len, bs, self.num_attention_heads, -1))
+        num_head = 1 if one_head else self.num_attention_heads
+        x = self.reshape(x, (seq_len, bs, num_head, -1))
         x = self.tnd_transpose(x, (1, 0, 2, 3))
-        x = self.reshape(x, (bs * seq_len, self.num_attention_heads, -1))
+        x = self.reshape(x, (bs * seq_len, num_head, -1))
         return x
 
     def _ulysses_qkv_a2a(self, qkv):
@@ -338,7 +404,6 @@ class MLASelfAttentionConcatenated(MultiLatentAttention):
         self.split_3d = SplitWithSize()
         self.tile_kv = Tile()
         self.pe_concat = Concat(axis=3)
-        self.pe_tnd_concat = Concat(axis=3)
         self.apply_rotary_emb = ApplyRotaryPosEmb(config)
         self.apply_rotary_emb2 = ApplyRotaryPosEmb(config, for_k_pos_emb=True)
         self.reshape = Reshape()
@@ -391,7 +456,7 @@ class MLASelfAttentionConcatenated(MultiLatentAttention):
         self.linear_kvb = build_module(
             submodules.linear_kvb,
             self.kv_lora_rank,
-            self.config.num_attention_heads * (self.q_head_dim - self.qk_pos_emb_head_dim + self.v_head_dim),
+            self.config.num_attention_heads * (self.qk_head_dim + self.v_head_dim),
             config=self.config,
             init_method=self.config.init_method,
             bias=self.config.add_bias_linear or self.config.add_qkv_bias,
@@ -479,16 +544,6 @@ class MLASelfAttentionConcatenated(MultiLatentAttention):
         k_pe = self.reshape(k_pe, (seq_len, bs, 1, self.qk_pos_emb_head_dim))
         compressed_kv_norm = self.k_layernorm(compressed_kv)
 
-        kv = self.linear_kvb(compressed_kv_norm)[0]
-        kv = self.reshape(kv, (
-            seq_len,
-            bs,
-            self.num_attention_heads,
-            self.qk_head_dim + self.v_head_dim,
-        ))
-
-        k_nope, value = self.split(kv, [self.qk_head_dim, self.v_head_dim], dim=-1)
-
         if rotary_pos_emb is not None:
             q_pe = self.apply_rotary_emb(
                 q_pe,
@@ -503,10 +558,27 @@ class MLASelfAttentionConcatenated(MultiLatentAttention):
                 multi_latent_attention=self.config.multi_latent_attention
             )
 
-        query = self.pe_concat([q_nope, q_pe])
-        k_pe = self.tile_kv(k_pe, (1, 1, self.num_attention_heads, 1))
-        key = self.pe_concat([k_nope, k_pe])
-
+        if self.use_mla_absorb:
+            k_nope = self.reshape(compressed_kv_norm, (seq_len, bs, 1, -1))
+            value = self.reshape(compressed_kv_norm, (seq_len, bs, 1, -1))
+            query = (q_nope, q_pe)
+            key = (k_nope, k_pe)
+        else:
+            kv = self.linear_kvb(compressed_kv_norm)[0]
+            kv = self.reshape(kv, (
+                seq_len,
+                bs,
+                self.num_attention_heads,
+                self.qk_head_dim + self.v_head_dim,
+            ))
+
+            k_nope, value = self.split(kv, [self.qk_head_dim, self.v_head_dim], dim=-1)
+            query = self.pe_concat([q_nope, q_pe])
+            k_pe = self.tile_kv(k_pe, (1, 1, self.num_attention_heads, 1))
+            key = self.pe_concat([k_nope, k_pe])
+
+        if self.use_dsa:
+            return query, key, value, q_a
         return query, key, value
 
 
@@ -538,7 +610,6 @@ class MLASelfAttention(MultiLatentAttention):
         self.split_3d = SplitWithSize()
         self.tile_kv = Tile()
         self.pe_concat = Concat(axis=3)
-        self.pe_tnd_concat = Concat(axis=3)
         self.apply_rotary_emb = ApplyRotaryPosEmb(config)
         self.apply_rotary_emb2 = ApplyRotaryPosEmb(config, for_k_pos_emb=True)
         self.reshape = Reshape()
diff --git a/mindformers/parallel_core/training_graph/transformer/transformer_block.py b/mindformers/parallel_core/training_graph/transformer/transformer_block.py
index 53c16acddf2d662bff1e1deb4afb443b2d6a7a72..165e0b72222bd9c1c3ce4612ef3880743c88e0fb 100644
--- a/mindformers/parallel_core/training_graph/transformer/transformer_block.py
+++ b/mindformers/parallel_core/training_graph/transformer/transformer_block.py
@@ -127,6 +127,7 @@ class TransformerBlock(nn.Cell):
         self.submodules = _get_block_submodules(config, spec)
         self.post_layer_norm = post_layer_norm
         self.num_layers = config.num_layers
+        self.use_dsa = config.experimental_attention_variant == "dsa"
         cp = 1 if config is None else config.context_parallel_size
         if config.sequence_parallel and cp > 1:
             logger.warning("The context paralley way conflicts with sequence with sequence parallel way."
@@ -144,6 +145,8 @@ class TransformerBlock(nn.Cell):
         self._build_layers(config)
 
         self.init_extra_loss = Tensor([0], mstype.float32)
+        if self.use_dsa:
+            self.init_indexer_loss = Tensor([0], mstype.float32)
 
         self.shard(config)
 
@@ -180,6 +183,8 @@ class TransformerBlock(nn.Cell):
                   actual_seq_len=None):
         """ Construct function of transformer. """
         extra_loss = self.init_extra_loss
+        if self.use_dsa:
+            indexer_loss = self.init_indexer_loss
         sharded_key = None
         sharded_value = None
         for index in range(self.num_layers):
@@ -202,6 +207,10 @@ class TransformerBlock(nn.Cell):
                 # has no practical impact on training.
                 # attention_bias is useless by default, only used for T5 in Megatron v0.12.0.
             )
+            # if self.use_dsa:
+            #     hidden_states, context, extra_loss, indexer_loss = layer_output
+            # else:
+            #     hidden_states, context, extra_loss = layer_output
 
         # final layernorm.
         if self.post_layer_norm:
diff --git a/mindformers/parallel_core/training_graph/transformer/transformer_layer.py b/mindformers/parallel_core/training_graph/transformer/transformer_layer.py
index a53f58664f35a68bd69ee0a686aec1bfa2e5b48f..1b944339b19ed75cd9088e0ee9a6631744db950f 100644
--- a/mindformers/parallel_core/training_graph/transformer/transformer_layer.py
+++ b/mindformers/parallel_core/training_graph/transformer/transformer_layer.py
@@ -95,6 +95,7 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer):
         self.sequence_parallel = config.sequence_parallel
         self.fp32_residual_connection = config.fp32_residual_connection
         self.compute_dtype = config.compute_dtype
+        self.use_dsa = config.experimental_attention_variant == "dsa"
 
         self.input_layernorm = build_module(
             submodules.input_layernorm,
@@ -217,6 +218,8 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer):
             prefix_keys_values=prefix_keys_values,
             actual_seq_len=actual_seq_len
         )
+        # if self.use_dsa:
+        #     attention_output_with_bias, indexer_loss = attention_output_with_bias
 
         if isinstance(attention_output_with_bias, tuple):
             attention_output, self_attention_bias = attention_output_with_bias
@@ -305,6 +308,8 @@ class TransformerLayer(nn.Cell, BaseTransformerLayer):
         output = self.cast(output, self.compute_dtype)
 
         # 'return context' is useless, this param may be deprecated later.
+        if self.use_dsa:
+            return output, context, extra_loss, indexer_loss
         return output, context, extra_loss
 
     def shard(self, config: TransformerConfig):
diff --git a/mindformers/parallel_core/transformer_config.py b/mindformers/parallel_core/transformer_config.py
index 3c5f29c206a4b3e24cfa373b4cbcda6faf0862e0..1d39ad1a52d032c4cabab55828dd2776ab4cd28b 100644
--- a/mindformers/parallel_core/transformer_config.py
+++ b/mindformers/parallel_core/transformer_config.py
@@ -118,6 +118,32 @@ class TransformerConfig(ModelParallelConfig, MFModelConfig):
     qk_layernorm: bool = False
     """Whether to apply `normalization` type of normalization to the query and key embeddings."""
 
+    ####################
+    # attention variant
+    ####################
+    experimental_attention_variant: Optional[str] = None
+    """
+    Type of attention variant to use. Currently support ['dsa'].
+    Defaults: `None`, which means that no attention variant will be applied.
+    """
+
+    dsa_indexer_n_heads: Optional[int] = None
+    """Number of DSA indexer heads."""
+
+    dsa_indexer_head_dim: Optional[int] = None
+    """Dimension per DSA indexer head."""
+
+    dsa_indexer_topk: Optional[int] = None
+    """Number of top-k tokens to select in DSA indexer."""
+
+    dsa_indexer_loss_coeff: Optional[float] = None
+    """Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss."""
+
+    dsa_indexer_use_sparse_loss: bool = False
+    """
+    Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the top-k indices.
+    """
+
     ####################
     # Initialization
     ####################
@@ -761,6 +787,14 @@ class TransformerConfig(ModelParallelConfig, MFModelConfig):
                              f"should be equal to num_layers. But get num_encoder_layers: {self.num_encoder_layers}, "
                              f"num_decoder_layers: {self.num_decoder_layers}, num_layers: {self.num_layers}.")
 
+        if self.experimental_attention_variant == "dsa":
+            if not self.dsa_indexer_use_sparse_loss and self.input_layout == "TND":
+                raise ValueError("Currently dense indexer loss is not supported in TND layout of attention!")
+            # assert (
+            #     self.context_parallel_size == 1
+            # ), "Currently context parallelism is not supported by DSAttention when!"
+            assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention!"
+
 
     def _validate_param_init_std_rules(self):
         """Validate and compile decoupling initialization rules."""
@@ -825,6 +859,11 @@ class MLATransformerConfig(TransformerConfig):
     multi_latent_attention: bool = True
     """Whether to use Multi-Latent Attention."""
 
+    use_mla_absorb: bool = False
+    """
+    Whether to use mla-absorb (MQA).
+    """
+
     q_lora_rank: int = 512
     """Rank of Query tensor's low rank representation."""