diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 1b0cf16315c..7ac46d42e05 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -2,6 +2,7 @@
 
 import logging
 import math
+import operator
 import warnings
 from contextlib import nullcontext
 from typing import List, Optional, Sequence, Tuple
@@ -29,7 +30,10 @@
 from megatron.core.inference.utils import device_memory_summary, tensor_swap
 from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
 from megatron.core.package_info import __version__ as mcore_version
-from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list
+from megatron.core.ssm.mamba_hybrid_layer_allocation import (
+    Symbols,
+    get_layer_maps_from_layer_type_list,
+)
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
 from megatron.core.utils import deprecate_args
 from megatron.core.utils import divide as core_divide
@@ -332,16 +336,18 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
             # For hybrid models, the layer map converts the global layer index to the
             # corresponding attention layer index or Mamba layer index depending on the
             # layer type.
-            mamba_layer_map, gdn_layer_map, attention_layer_map, _, _ = (
-                get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)
+            attention_layer_map, dsa_layer_map, gdn_layer_map, mamba_layer_map = (
+                operator.itemgetter(
+                    Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.MAMBA
+                )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list))
             )
 
             if len(gdn_layer_map) > 0:
                 raise NotImplementedError("GDN layers are not supported for inference.")
 
-            self.num_attention_layers = len(attention_layer_map)
+            self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map)
             self.num_mamba_layers = len(mamba_layer_map)
-            self.layer_map = attention_layer_map | mamba_layer_map
+            self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map
         else:
             # The layer map is the identity function for pure Transformer models.
             self.num_attention_layers = model_config.num_layers // pp_size
diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
index 6608073136c..1b03b935639 100644
--- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
+++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
@@ -123,6 +123,7 @@ def get_dsa_module_spec_for_backend(
             q_layernorm=IdentityOp,
             kv_layernorm=IdentityOp,
         ),
+        metainfo={"fuse_input_layernorm": False},
     )
 
     return attention
@@ -138,6 +139,8 @@ def get_experimental_attention_variant_module_spec(
 
     if config.experimental_attention_variant == "gated_delta_net":
         return get_gated_delta_net_module_spec(config=config, backend=backend)
+    elif config.experimental_attention_variant == "dsa":
+        return get_dsa_module_spec_for_backend(config=config, backend=backend)
     else:
         raise ValueError(
             f"Invalid experimental attention variant: {config.experimental_attention_variant}"
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
index 48f25bdbab9..96b2e13e2d4 100755
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -4,6 +4,7 @@
     TEColumnParallelLinear,
     TEDotProductAttention,
     TELayerNormColumnParallelLinear,
+    TELinear,
     TENorm,
     TERowParallelLinear,
 )
@@ -24,7 +25,18 @@
 )
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant.dsa import (
+    DSAIndexer,
+    DSAIndexerSubmodules,
+    DSAttention,
+    DSAttentionSubmodules,
+)
+from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.multi_latent_attention import (
+    MLASelfAttention,
+    MLASelfAttentionSubmodules,
+)
 from megatron.core.transformer.multi_token_prediction import (
     MultiTokenPredictionBlock,
     MultiTokenPredictionBlockSubmodules,
@@ -117,6 +129,41 @@
                 self_attn_bda=get_bias_dropout_add,
             ),
         ),
+        dsa_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=TENorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=TEColumnParallelLinear,
+                        linear_q_down_proj=TELinear,
+                        linear_q_up_proj=TEColumnParallelLinear,
+                        linear_kv_down_proj=TELinear,
+                        linear_kv_up_proj=TEColumnParallelLinear,
+                        core_attention=ModuleSpec(
+                            module=DSAttention,
+                            submodules=DSAttentionSubmodules(
+                                indexer=ModuleSpec(
+                                    module=DSAIndexer,
+                                    submodules=DSAIndexerSubmodules(
+                                        linear_wq_b=TELinear,
+                                        linear_wk=TELinear,
+                                        k_norm=TENorm,
+                                        linear_weights_proj=TELinear,
+                                    ),
+                                )
+                            ),
+                        ),
+                        linear_proj=TERowParallelLinear,
+                        q_layernorm=IdentityOp,
+                        kv_layernorm=IdentityOp,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
         # Started with spec from gpt_layer_specs.py
         # Using the TE spec because we had problems getting the non-TE spec
         # working
@@ -177,6 +224,41 @@
                 self_attn_bda=get_bias_dropout_add,
             ),
         ),
+        dsa_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=TENorm,
+                self_attention=ModuleSpec(
+                    module=MLASelfAttention,
+                    params={"attn_mask_type": AttnMaskType.causal},
+                    submodules=MLASelfAttentionSubmodules(
+                        linear_q_proj=TEColumnParallelLinear,
+                        linear_q_down_proj=TELinear,
+                        linear_q_up_proj=TEColumnParallelLinear,
+                        linear_kv_down_proj=TELinear,
+                        linear_kv_up_proj=TEColumnParallelLinear,
+                        core_attention=ModuleSpec(
+                            module=DSAttention,
+                            submodules=DSAttentionSubmodules(
+                                indexer=ModuleSpec(
+                                    module=DSAIndexer,
+                                    submodules=DSAIndexerSubmodules(
+                                        linear_wq_b=TELinear,
+                                        linear_wk=TELinear,
+                                        k_norm=TENorm,
+                                        linear_weights_proj=TELinear,
+                                    ),
+                                )
+                            ),
+                        ),
+                        linear_proj=InferenceRowParallelLinear,
+                        q_layernorm=IdentityOp,
+                        kv_layernorm=IdentityOp,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
         # Started with spec from gpt_layer_specs.py
         # Using the TE spec because we had problems getting the non-TE spec
         # working
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 817db8789d3..d0458988dcc 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -217,7 +217,9 @@ def __init__(
                 tp_group=self.pg_collection.tp,
             )
 
-        if self.position_embedding_type == 'rope':
+        # MLA (also used by DeepSeek Sparse Attention) uses its own decoupled RoPE, therefore we do
+        # not build standard RoPE here when using MLA.
+        if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
             self.rotary_pos_emb = RotaryEmbedding(
                 kv_channels=self.config.kv_channels,
                 rotary_percent=rotary_percent,
@@ -373,7 +375,7 @@ def forward(
             decoder_input = None
 
         rotary_pos_emb = None
-        if self.position_embedding_type == 'rope':
+        if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
             rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
                 inference_context, self.decoder, decoder_input, self.config, packed_seq_params
             )
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
index f42f3542c3d..fa262b2293f 100644
--- a/megatron/core/ssm/mamba_block.py
+++ b/megatron/core/ssm/mamba_block.py
@@ -41,6 +41,7 @@ class MambaStackSubmodules:
     mamba_layer: Union[ModuleSpec, type] = IdentityOp
     gdn_layer: Union[ModuleSpec, type] = IdentityOp
     attention_layer: Union[ModuleSpec, type] = IdentityOp
+    dsa_layer: Union[ModuleSpec, type] = IdentityOp
     mlp_layer: Union[ModuleSpec, type] = IdentityOp
     moe_layer: Union[ModuleSpec, type] = IdentityOp
     mtp_block_spec: Optional[ModuleSpec] = None
@@ -135,6 +136,16 @@ def __init__(
                         add_layer_offset=False,
                         pp_layer_offset=pp_layer_offset,
                     )
+                elif layer_type == LayerSymbols.DS_ATTENTION:
+                    layer = build_module(
+                        submodules.dsa_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
+                        add_layer_offset=False,
+                        pp_layer_offset=pp_layer_offset,
+                    )
                 elif layer_type == LayerSymbols.MLP:
                     layer = build_module(
                         submodules.mlp_layer,
@@ -161,7 +172,7 @@ def __init__(
                         add_layer_offset=False,
                     )
                 else:
-                    assert False, "unexpected layer_type"
+                    raise ValueError("unexpected layer_type")
             self.layers.append(layer)
 
         # Required for activation recomputation
diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
index 1cb1e4a31d1..1947dbc39f1 100644
--- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py
+++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
@@ -17,11 +17,24 @@ class Symbols:
     MAMBA = "M"
     GDN = 'G'
     ATTENTION = "*"
+    DS_ATTENTION = "D"
     MLP = "-"
     MOE = 'E'
     PIPE = '|'
     MTP_SEPARATOR = "/"
-    VALID_LAYERS = {MAMBA, GDN, ATTENTION, MLP, MOE}
+    VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLP, MOE}
+
+    @classmethod
+    def name_sorted_valid_layer_symbols(cls) -> list[str]:
+        """Return the valid layer symbols sorted lexicographically by their public attribute
+        name.
+        """
+        valid_layer_attrs = []
+        for name, value in vars(cls).items():
+            if not name.startswith('_') and value in cls.VALID_LAYERS:
+                valid_layer_attrs.append((name, value))
+        valid_layer_attrs.sort()
+        return [value for (_, value) in valid_layer_attrs]
 
 
 @dataclass
@@ -155,24 +168,18 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]:
         pattern: Full hybrid layer pattern string.
 
     Returns:
-        Dictionary mapping layer symbol to count. Keys are Symbols.MAMBA,
-        Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, and Symbols.MOE.
+        Dictionary mapping layer symbol to count. Keys are all valid layer symbols
+            (Symbols.VALID_LAYERS).
 
     Examples:
         >>> get_hybrid_layer_counts("M*M*")
-        {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0}
+        {'*': 2, 'G': 0, 'D': 0, 'M': 2, '-': 0, 'E': 0}
 
         >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM")
-        {'M': 8, 'G': 0, '*': 1, '-': 4, 'E': 0}
+        {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0}
     """
     parsed = parse_hybrid_pattern(pattern)
-    counts = {
-        Symbols.MAMBA: 0,
-        Symbols.GDN: 0,
-        Symbols.ATTENTION: 0,
-        Symbols.MLP: 0,
-        Symbols.MOE: 0,
-    }
+    counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()}
 
     # Count main decoder layers (skip '|' pipe separators)
     if parsed.main_pattern:
@@ -285,6 +292,10 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False)
                 f"Valid symbols are: {valid_chars}"
             )
 
+    # Disallow Attention + MLA/DSA hybridity.
+    if Symbols.ATTENTION in pattern and Symbols.DS_ATTENTION in pattern:
+        raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+
 
 def validate_segment_layers(segment: str) -> List[str]:
     """Validate and convert a single pipeline segment pattern to a layer type list.
@@ -308,6 +319,10 @@ def validate_segment_layers(segment: str) -> List[str]:
                 f"In hybrid layer pattern segment, '{layer_char}' is not "
                 f"one of {Symbols.VALID_LAYERS}"
             )
+
+    # Disallow Attention + MLA/DSA hybridity.
+    if Symbols.ATTENTION in segment and Symbols.DS_ATTENTION in segment:
+        raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
     return layer_type_list
 
 
@@ -468,17 +483,15 @@ def select_pipeline_segment(
     return layer_type_list, layer_offset
 
 
-def get_layer_maps_from_layer_type_list(
-    layer_type_list: List[str],
-) -> Tuple[Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int]]:
+def get_layer_maps_from_layer_type_list(layer_type_list: list[str]) -> dict[str, dict[int, int]]:
     """
     Returns maps from global layer index to the corresponding layer index
-    for each layer type in [Mamba, GDN, Attention, MLP, MoE] given a layer type list.
+    for each valid layer type (those in Symbols.VALID_LAYERS) given a layer type list.
     """
-    layer_types = [Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE]
+    layer_types = [symbol for symbol in Symbols.name_sorted_valid_layer_symbols()]
     layer_maps = {layer_type: {} for layer_type in layer_types}
     for global_layer_idx, layer_type in enumerate(layer_type_list):
         layer_map = layer_maps[layer_type]
         local_layer_idx = len(layer_map)
         layer_map[global_layer_idx] = local_layer_idx
-    return [layer_maps[layer_type] for layer_type in layer_types]
+    return layer_maps
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index 605cb585779..0e5cec991e8 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -867,7 +867,9 @@ def flash_decode_and_prefill(
             q = q.reshape(num_requests, tokens_per_request, q.shape[2], q.shape[3])
 
             # If using MLA we use the FlashMLA kernel
-            if isinstance(self.config, MLATransformerConfig):
+            # The `softmax_scale` attribute check is to find out whether this is an MLA layer or
+            # standard Attention.
+            if isinstance(self.config, MLATransformerConfig) and hasattr(self, "softmax_scale"):
                 softmax_scale = self.softmax_scale
 
                 num_heads_k = 1  # Only a single head for MLA Flash
diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
index a27fed14d6b..601ae89fae1 100644
--- a/megatron/core/transformer/multi_latent_attention.py
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -138,6 +138,7 @@ def __init__(
         attention_type: str,
         cp_comm_type: Optional[str] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        pp_layer_offset: Optional[int] = None,
     ) -> None:
 
         super().__init__(
@@ -147,6 +148,7 @@ def __init__(
             attention_type=attention_type,
             attn_mask_type=attn_mask_type,
             pg_collection=pg_collection,
+            pp_layer_offset=pp_layer_offset,
         )
         self.config: MLATransformerConfig
 
@@ -473,6 +475,7 @@ def __init__(
         attn_mask_type=AttnMaskType.padding,
         cp_comm_type: Optional[str] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        pp_layer_offset: Optional[int] = None,
     ):
         if pg_collection is None:
             pg_collection = ProcessGroupCollection.use_mpu_process_groups()
@@ -485,6 +488,7 @@ def __init__(
             attention_type="self",
             cp_comm_type=cp_comm_type,
             pg_collection=pg_collection,
+            pp_layer_offset=pp_layer_offset,
         )
 
         if self.config.q_lora_rank is None:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 31433704b6b..0413ea518c9 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -799,6 +799,11 @@ def validate_args(args, defaults={}):
                 "This argument will be ignored.",
                 args.rank
             )
+
+    # Infer use of MLA from unified pattern
+    if args.hybrid_layer_pattern and Symbols.DS_ATTENTION in args.hybrid_layer_pattern:
+        args.multi_latent_attention = True
+
     # === End of hybrid layer pattern: deprecation handling and validation ===
 
     # Uneven virtual pipeline parallelism
@@ -1763,6 +1768,9 @@ def core_transformer_config_from_args(args, config_class=None):
         kw_args['cp_comm_type'] = args.cp_comm_type[0]
     if args.hybrid_layer_pattern is not None:
         kw_args['is_hybrid_model'] = True
+        from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
+        if Symbols.DS_ATTENTION in args.hybrid_layer_pattern:
+            kw_args['experimental_attention_variant'] = 'dsa'
 
     kw_args['inference_sampling_seed'] = args.seed
 
@@ -3286,9 +3294,10 @@ def _add_experimental_args(parser):
                        '`transformer_block.py`, or `transformer_layer.py`')
     group.add_argument('--hybrid-layer-pattern', type=str, default=None,
                        help='Specify a hybrid layer pattern using M (mamba), G (gdn), '
-                       '* (attention), - (mlp), E (moe). Use | to define pipeline stage '
-                       'boundaries for flexible virtual pipeline parallel (fVPP). Use / to '
-                       'separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". '
+                       '* (attention), D (dsa), - (mlp), E (moe). Use | to define pipeline '
+                       'stage boundaries for flexible virtual pipeline parallel (fVPP). '
+                       'Use / to separate MTP patterns. '
+                       'Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". '
                        'When this flag is used, it is the sole indicator that a hybrid model '
                        'is being run.')
     group.add_argument('--hybrid-override-pattern', type=str, default=None,
diff --git a/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py b/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py
new file mode 100644
index 00000000000..96b782fad85
--- /dev/null
+++ b/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py
@@ -0,0 +1,644 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+"""
+Equivalence tests: GPTModel with DSA vs MambaModel with DSA pattern.
+
+A small DeepSeek-V3.2 proxy model (4 GPT layers / 8 Mamba layers) is built,
+weights are remapped GPT→Mamba, and logprobs are compared to verify they are
+numerically identical.
+
+Architecture equivalence
+------------------------
+GPTModel layer N  (combined attention + MLP in one TransformerLayer)
+  ≡  MambaModel layer 2N  (D, DSA TransformerLayer: input_layernorm + MLASelfAttention)
+   + MambaModel layer 2N+1 (-, MLPLayer: fused-norm MLP)
+
+Run with::
+
+    torchrun --nproc-per-node=2 -m pytest \\
+        tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py -v
+"""
+
+import copy
+import json
+import math
+import os
+from pathlib import Path
+from typing import Dict, Optional
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+
+import megatron.core.parallel_state as mpu
+from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+    get_transformer_block_with_experimental_attention_variant_spec,
+)
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec
+from megatron.core.models.mamba.mamba_model import MambaModel
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.ssm.mamba_hybrid_layer_allocation import validate_segment_layers
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import MLATransformerConfig
+from megatron.rl.rl_utils import selective_log_softmax
+from tests.unit_tests.test_utilities import Utils
+
+try:
+    from fast_hadamard_transform import hadamard_transform as _hadamard_transform
+
+    HAVE_HADAMARD = True
+except ImportError:
+    HAVE_HADAMARD = False
+
+
+# ---------------------------------------------------------------------------
+# Hadamard mock (used when the library is not installed)
+# ---------------------------------------------------------------------------
+
+
+def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+    """Identity-scale mock for hadamard_transform used in DSA."""
+    return x * scale
+
+
+@pytest.fixture(autouse=True)
+def _patch_hadamard_if_needed():
+    """Patch hadamard_transform in the DSA module when the library is absent."""
+    if not HAVE_HADAMARD:
+        with patch(
+            'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform',
+            _mock_hadamard_transform,
+        ):
+            yield
+    else:
+        yield
+
+
+# ---------------------------------------------------------------------------
+# Proxy model constants
+# ---------------------------------------------------------------------------
+
+_VOCAB_SIZE = 256
+_MAX_SEQ_LEN = 64
+_SEQ_LEN = 32
+_BATCH_SIZE = 2
+_NUM_GPT_LAYERS = 4
+_MAMBA_PATTERN = "D-D-D-D-"  # len=8 = 2 * _NUM_GPT_LAYERS
+
+# MoE variant: first 2 GPT layers dense, last 2 MoE → 8 Mamba layers
+_MOE_MAMBA_PATTERN = "D-D-DEDE"  # 2 dense (D-) + 2 MoE (DE)
+
+
+# ---------------------------------------------------------------------------
+# Model construction helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_dsa_config(num_layers: int, tp: int = 1, pp: int = 1) -> MLATransformerConfig:
+    """Return a small DeepSeek-V3.2 proxy MLATransformerConfig."""
+    return MLATransformerConfig(
+        num_layers=num_layers,
+        hidden_size=256,
+        num_attention_heads=16,
+        q_lora_rank=64,
+        kv_lora_rank=64,
+        qk_head_dim=64,
+        qk_pos_emb_head_dim=32,
+        v_head_dim=64,
+        dsa_indexer_n_heads=8,
+        dsa_indexer_head_dim=64,
+        dsa_indexer_topk=32,
+        normalization="RMSNorm",
+        bf16=True,
+        params_dtype=torch.bfloat16,
+        pipeline_dtype=torch.bfloat16,
+        add_bias_linear=False,
+        use_cpu_initialization=True,
+        rope_type='rope',
+        experimental_attention_variant="dsa",
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        tensor_model_parallel_size=tp,
+        pipeline_model_parallel_size=pp,
+    )
+
+
+def _make_dsa_moe_config(num_layers: int, tp: int = 1, pp: int = 1) -> MLATransformerConfig:
+    """Return a small DeepSeek-V3 proxy MLATransformerConfig with MoE layers.
+
+    Mirrors the DeepSeek-V3 pattern: first 2 GPT layers are dense, last 2 are MoE.
+    ``moe_layer_freq=[0, 0, 1, 1]`` controls which GPT layers become MoE layers.
+    """
+    return MLATransformerConfig(
+        num_layers=num_layers,
+        hidden_size=256,
+        num_attention_heads=16,
+        q_lora_rank=64,
+        kv_lora_rank=64,
+        qk_head_dim=64,
+        qk_pos_emb_head_dim=32,
+        v_head_dim=64,
+        dsa_indexer_n_heads=8,
+        dsa_indexer_head_dim=64,
+        dsa_indexer_topk=32,
+        normalization="RMSNorm",
+        bf16=True,
+        params_dtype=torch.bfloat16,
+        pipeline_dtype=torch.bfloat16,
+        add_bias_linear=False,
+        use_cpu_initialization=True,
+        rope_type='rope',
+        experimental_attention_variant="dsa",
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        tensor_model_parallel_size=tp,
+        # Enable sequence parallelism when TP is used, otherwise we
+        # error out due to use of MoE.
+        sequence_parallel=tp > 1,
+        pipeline_model_parallel_size=pp,
+        # MoE fields
+        num_moe_experts=4,
+        moe_router_topk=2,
+        moe_grouped_gemm=True,
+        moe_token_dispatcher_type="allgather",
+        moe_router_load_balancing_type="aux_loss",
+        moe_aux_loss_coeff=0.0,
+        moe_ffn_hidden_size=512,
+        moe_shared_expert_intermediate_size=512,
+        moe_layer_freq=[0, 0, 1, 1],  # first 2 layers dense, last 2 MoE
+    )
+
+
+def _build_gpt_model(
+    config: MLATransformerConfig, pre_process: bool = True, post_process: bool = True
+) -> GPTModel:
+    """Build a GPTModel with the DSA transformer block spec."""
+    spec = get_transformer_block_with_experimental_attention_variant_spec(config)
+    model = GPTModel(
+        config=config,
+        transformer_layer_spec=spec,
+        vocab_size=_VOCAB_SIZE,
+        max_sequence_length=_MAX_SEQ_LEN,
+        pre_process=pre_process,
+        post_process=post_process,
+        parallel_output=False,  # Gather logits across TP for easy comparison
+        position_embedding_type='rope',
+    )
+    return model.cuda()
+
+
+def _build_mamba_model(
+    config: MLATransformerConfig,
+    layer_pattern: str,
+    pre_process: bool = True,
+    post_process: bool = True,
+) -> MambaModel:
+    """Build a MambaModel with the given hybrid layer pattern."""
+    layer_type_list = validate_segment_layers(layer_pattern)
+    mamba_config = copy.deepcopy(config)
+    mamba_config.num_layers = len(layer_type_list)
+    assert mamba_config.num_layers == _NUM_GPT_LAYERS * 2
+    model = MambaModel(
+        config=mamba_config,
+        mamba_stack_spec=mamba_stack_spec,
+        vocab_size=_VOCAB_SIZE,
+        max_sequence_length=_MAX_SEQ_LEN,
+        pre_process=pre_process,
+        post_process=post_process,
+        parallel_output=False,
+        hybrid_layer_pattern=layer_pattern,
+        position_embedding_type='rope',
+    )
+    return model.cuda()
+
+
+# ---------------------------------------------------------------------------
+# Weight remapping
+# ---------------------------------------------------------------------------
+
+
+def _remap_gpt_to_mamba_state_dict(
+    gpt_sd: Dict[str, torch.Tensor], num_local_gpt_layers: int
+) -> Dict[str, torch.Tensor]:
+    """Remap a GPTModel state_dict to a MambaModel state_dict.
+
+    GPTModel layer N (combined attention + MLP) maps to:
+      * MambaModel layer 2N   – DSA attention (input_layernorm + self_attention)
+      * MambaModel layer 2N+1 – MLP           (mlp.*)
+
+    Additionally, ``decoder.final_layernorm.*`` (TransformerBlock naming) is
+    remapped to ``decoder.final_norm.*`` (MambaStack naming).
+
+    All other keys (embedding, output_layer, rotary_pos_emb, …) are unchanged.
+
+    Args:
+        gpt_sd: State dict obtained from GPTModel.state_dict().
+        num_local_gpt_layers: Number of GPT decoder layers on the current
+            pipeline stage (i.e. ``len(gpt_model.decoder.layers)``).
+
+    Returns:
+        Remapped state dict ready for MambaModel.load_state_dict(strict=True).
+    """
+    mamba_sd: Dict[str, torch.Tensor] = {}
+    layer_prefix = "decoder.layers."
+    final_ln_prefix = "decoder.final_layernorm."
+
+    for key, value in gpt_sd.items():
+        # ---- final layernorm rename ----
+        if key.startswith(final_ln_prefix):
+            suffix = key[len(final_ln_prefix) :]
+            mamba_sd[f"decoder.final_norm.{suffix}"] = value
+            continue
+
+        # ---- non-layer keys pass through unchanged ----
+        if not key.startswith(layer_prefix):
+            mamba_sd[key] = value
+            continue
+
+        # ---- parse "decoder.layers.{N}.{rest}" ----
+        remainder = key[len(layer_prefix) :]
+        dot_idx = remainder.index('.')
+        layer_n = int(remainder[:dot_idx])
+        rest = remainder[dot_idx + 1 :]  # e.g. "self_attention.linear_q_proj.weight"
+
+        assert (
+            0 <= layer_n < num_local_gpt_layers
+        ), f"Layer index {layer_n} out of range [0, {num_local_gpt_layers}) in key '{key}'"
+
+        if rest.startswith("input_layernorm.") or rest.startswith("self_attention."):
+            # Attention sub-module → DSA layer 2N
+            mamba_sd[f"{layer_prefix}{2 * layer_n}.{rest}"] = value
+        elif rest.startswith("mlp."):
+            # MLP sub-module → MLP layer 2N+1
+            mamba_sd[f"{layer_prefix}{2 * layer_n + 1}.{rest}"] = value
+        elif rest.startswith("pre_mlp_layernorm."):
+            # pre_mlp_layernorm → MoE layer 2N+1 (MoETransformerLayer has TENorm)
+            # Dense layers use IdentityOp for pre_mlp_layernorm (no state dict keys),
+            # so this branch only fires for MoE layers.
+            mamba_sd[f"{layer_prefix}{2 * layer_n + 1}.{rest}"] = value
+        else:
+            # self_attn_bda / mlp_bda are callables (no weights).
+            # Anything else is unexpected.
+            raise ValueError(
+                f"Unexpected sub-key '{rest}' in GPT layer {layer_n} (full key='{key}'). "
+                "Expected: input_layernorm.*, self_attention.*, pre_mlp_layernorm.*, mlp.*"
+            )
+
+    return mamba_sd
+
+
+# ---------------------------------------------------------------------------
+# Forward-pass helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_inputs(tokens: torch.Tensor):
+    """Return position_ids and attention_mask for a token batch."""
+    batch_size, seq_len = tokens.shape
+    position_ids = (
+        torch.arange(seq_len, device=tokens.device).unsqueeze(0).expand(batch_size, seq_len)
+    )
+    attention_mask = torch.ones(
+        batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=tokens.device
+    )
+    return position_ids, attention_mask
+
+
+def _forward_logprobs_pp1(model: torch.nn.Module, tokens: torch.Tensor) -> torch.Tensor:
+    """Single-stage (PP=1) forward returning logprobs [batch, seq-1]."""
+    position_ids, attention_mask = _make_inputs(tokens)
+    with torch.no_grad():
+        logits = model(input_ids=tokens, position_ids=position_ids, attention_mask=attention_mask)
+    return selective_log_softmax(logits[:, :-1, :], tokens[:, 1:])
+
+
+def _forward_logprobs_pp2(model: torch.nn.Module, tokens: torch.Tensor) -> Optional[torch.Tensor]:
+    """Two-stage (PP=2) forward using point-to-point communication.
+
+    Returns logprobs on the last PP stage; None on the first stage.
+    The caller must invoke this function for both GPT and Mamba models in the
+    *same order* on all ranks to avoid deadlocks.
+    """
+    batch_size, seq_len = tokens.shape
+    hidden_size = model.config.hidden_size
+    position_ids, attention_mask = _make_inputs(tokens)
+
+    pp_rank = mpu.get_pipeline_model_parallel_rank()
+    next_rank = mpu.get_pipeline_model_parallel_next_rank()
+    prev_rank = mpu.get_pipeline_model_parallel_prev_rank()
+
+    if pp_rank == 0:
+        # First stage: embedding + local layers → hidden states [seq, batch, hidden]
+        with torch.no_grad():
+            hidden = model(
+                input_ids=tokens, position_ids=position_ids, attention_mask=attention_mask
+            )
+        dist.send(hidden.contiguous(), dst=next_rank)
+        return None
+    else:
+        # Last stage: receive hidden states, run remaining layers → logits
+        hidden_buf = torch.empty(
+            seq_len, batch_size, hidden_size, dtype=torch.bfloat16, device=tokens.device
+        )
+        dist.recv(hidden_buf, src=prev_rank)
+        model.set_input_tensor(hidden_buf)
+        with torch.no_grad():
+            logits = model(
+                input_ids=tokens, position_ids=position_ids, attention_mask=attention_mask
+            )
+        return selective_log_softmax(logits[:, :-1, :], tokens[:, 1:])
+
+
+# ---------------------------------------------------------------------------
+# Golden-value comparison
+# ---------------------------------------------------------------------------
+
+
+def _compare_against_golden_values(
+    logprobs: torch.Tensor, golden_logprobs: torch.Tensor, abs_tol: float = 1e-3
+):
+    """Assert logprobs match the golden values JSON within *abs_tol*."""
+    golden_lp = golden_logprobs[0].float().tolist()
+    actual_lp = logprobs[0].float().tolist()
+    assert len(actual_lp) == len(
+        golden_lp
+    ), f"Logprob length mismatch: actual={len(actual_lp)}, golden={len(golden_lp)}"
+    for i, (a, g) in enumerate(zip(actual_lp, golden_lp)):
+        assert math.isclose(
+            a, g, abs_tol=abs_tol
+        ), f"Logprob mismatch at position {i}: actual={a:.6f}, golden={g:.6f}"
+
+
+# ---------------------------------------------------------------------------
+# Test class
+# ---------------------------------------------------------------------------
+
+_GOLDEN_BASE = Path(__file__).parent.parent.parent / ("functional_tests/test_cases/hybrid")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("tp,pp", [(1, 1), (2, 1), (1, 2)])
+class TestDSAGPTMambaEquivalence:
+    """Verify logprob equivalence between GPTModel+DSA and MambaModel+DSA.
+
+    For each distributed configuration (TP, PP), the test:
+    1. Builds a GPTModel with 4 DSA layers.
+    2. Builds a MambaModel with pattern "D-D-D-D-" (8 layers).
+    3. Remaps and loads GPT weights into MambaModel (strict=True).
+    4. Runs the same random tokens through both models.
+    5. Asserts logprob tensors are numerically close.
+    """
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    def _skip_if_insufficient_gpus(self, tp: int, pp: int) -> None:
+        world_size = int(os.environ.get('WORLD_SIZE', '1'))
+        required = tp * pp
+        if world_size < required:
+            pytest.skip(
+                f"Test tp={tp} pp={pp} requires {required} GPU(s), " f"but WORLD_SIZE={world_size}"
+            )
+
+    def test_dsa_logprobs_match(self, tp: int, pp: int) -> None:
+        """Build both models, transfer weights, compare logprobs."""
+        self._skip_if_insufficient_gpus(tp, pp)
+        Utils.initialize_model_parallel(tp, pp)
+        model_parallel_cuda_manual_seed(42)
+
+        pre_process = mpu.is_pipeline_first_stage()
+        post_process = mpu.is_pipeline_last_stage()
+
+        # ---- Build GPTModel ----
+        gpt_config = _make_dsa_config(num_layers=_NUM_GPT_LAYERS, tp=tp, pp=pp)
+        gpt_model = _build_gpt_model(gpt_config, pre_process=pre_process, post_process=post_process)
+        num_local_gpt_layers = len(gpt_model.decoder.layers)
+        gpt_sd = gpt_model.state_dict()
+
+        # ---- Build MambaModel ----
+        mamba_model = _build_mamba_model(
+            gpt_config, _MAMBA_PATTERN, pre_process=pre_process, post_process=post_process
+        )
+
+        # ---- Remap GPT weights → Mamba ----
+        mamba_sd = _remap_gpt_to_mamba_state_dict(gpt_sd, num_local_gpt_layers)
+        missing, unexpected = mamba_model.load_state_dict(mamba_sd, strict=True)
+        assert not missing, f"Missing keys after weight remap: {missing}"
+        assert not unexpected, f"Unexpected keys after weight remap: {unexpected}"
+
+        # ---- Create identical inputs on all ranks ----
+        torch.manual_seed(99)
+        tokens = torch.randint(0, _VOCAB_SIZE, (_BATCH_SIZE, _SEQ_LEN), device='cuda')
+
+        # ---- Forward pass ----
+        if pp == 1:
+            gpt_logprobs = _forward_logprobs_pp1(gpt_model, tokens)
+            mamba_logprobs = _forward_logprobs_pp1(mamba_model, tokens)
+            # Both models have full logits; compare on all TP ranks
+            torch.testing.assert_close(
+                gpt_logprobs,
+                mamba_logprobs,
+                atol=1e-5,
+                rtol=1e-5,
+                msg=f"Logprob mismatch for tp={tp} pp={pp}",
+            )
+        else:
+            # PP=2: manual pipeline communication
+            # Run GPT then Mamba in the same order on all ranks to avoid deadlocks.
+            gpt_logprobs = _forward_logprobs_pp2(gpt_model, tokens)
+            mamba_logprobs = _forward_logprobs_pp2(mamba_model, tokens)
+            if post_process:
+                torch.testing.assert_close(
+                    gpt_logprobs,
+                    mamba_logprobs,
+                    atol=1e-5,
+                    rtol=1e-5,
+                    msg=f"Logprob mismatch for tp={tp} pp={pp}",
+                )
+
+    def test_weight_loading_strict(self, tp: int, pp: int) -> None:
+        """Verify that strict=True weight loading succeeds (no missing/unexpected keys)."""
+        self._skip_if_insufficient_gpus(tp, pp)
+        Utils.initialize_model_parallel(tp, pp)
+        model_parallel_cuda_manual_seed(42)
+
+        pre_process = mpu.is_pipeline_first_stage()
+        post_process = mpu.is_pipeline_last_stage()
+
+        gpt_config = _make_dsa_config(num_layers=_NUM_GPT_LAYERS, tp=tp, pp=pp)
+        gpt_model = _build_gpt_model(gpt_config, pre_process=pre_process, post_process=post_process)
+        mamba_model = _build_mamba_model(
+            gpt_config, _MAMBA_PATTERN, pre_process=pre_process, post_process=post_process
+        )
+
+        gpt_sd = gpt_model.state_dict()
+        num_local_gpt_layers = len(gpt_model.decoder.layers)
+        mamba_sd = _remap_gpt_to_mamba_state_dict(gpt_sd, num_local_gpt_layers)
+        missing, unexpected = mamba_model.load_state_dict(mamba_sd, strict=True)
+
+        assert not missing, f"Missing keys: {missing}"
+        assert not unexpected, f"Unexpected keys: {unexpected}"
+
+    def test_record_and_compare_golden_values(self, tp: int, pp: int) -> None:
+        """Record GPTModel logprobs as golden values, then compare MambaModel against them.
+
+        Golden values are written to the functional test directory so they can be
+        committed and used by the CI inference golden-value tests.
+        """
+        self._skip_if_insufficient_gpus(tp, pp)
+        # Only run for TP=1, PP=1 (the canonical golden-value configuration)
+        if tp != 1 or pp != 1:
+            pytest.skip("Golden-value recording only runs for tp=1, pp=1")
+
+        Utils.initialize_model_parallel(1, 1)
+        model_parallel_cuda_manual_seed(42)
+
+        gpt_config = _make_dsa_config(num_layers=_NUM_GPT_LAYERS, tp=1, pp=1)
+        gpt_model = _build_gpt_model(gpt_config)
+        mamba_model = _build_mamba_model(gpt_config, _MAMBA_PATTERN)
+
+        gpt_sd = gpt_model.state_dict()
+        mamba_sd = _remap_gpt_to_mamba_state_dict(gpt_sd, len(gpt_model.decoder.layers))
+        mamba_model.load_state_dict(mamba_sd, strict=True)
+
+        torch.manual_seed(99)
+        tokens = torch.randint(0, _VOCAB_SIZE, (_BATCH_SIZE, _SEQ_LEN), device='cuda')
+
+        gpt_logprobs = _forward_logprobs_pp1(gpt_model, tokens)
+        mamba_logprobs = _forward_logprobs_pp1(mamba_model, tokens)
+
+        # Verify MambaModel matches golden values
+        _compare_against_golden_values(mamba_logprobs, gpt_logprobs, abs_tol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# MoE test class
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("tp,pp", [(1, 1), (2, 1), (1, 2)])
+class TestDSAMoEGPTMambaEquivalence:
+    """Verify logprob equivalence between GPTModel+DSA+MoE and MambaModel+DSA+MoE.
+
+    Architecture: 4 GPT layers with moe_layer_freq=[0,0,1,1] (first 2 dense, last 2 MoE)
+    maps to 8 Mamba layers with pattern "D-D-DEDE":
+      GPT layer 0 (dense) → Mamba layers 0 (D) + 1 (-)
+      GPT layer 1 (dense) → Mamba layers 2 (D) + 3 (-)
+      GPT layer 2 (MoE)   → Mamba layers 4 (D) + 5 (E)
+      GPT layer 3 (MoE)   → Mamba layers 6 (D) + 7 (E)
+    """
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    def _skip_if_insufficient_gpus(self, tp: int, pp: int) -> None:
+        world_size = int(os.environ.get('WORLD_SIZE', '1'))
+        required = tp * pp
+        if world_size < required:
+            pytest.skip(
+                f"Test tp={tp} pp={pp} requires {required} GPU(s), but WORLD_SIZE={world_size}"
+            )
+
+    def test_dsa_moe_logprobs_match(self, tp: int, pp: int) -> None:
+        """Build both models with MoE, transfer weights, compare logprobs."""
+        self._skip_if_insufficient_gpus(tp, pp)
+        Utils.initialize_model_parallel(tp, pp)
+        model_parallel_cuda_manual_seed(42)
+
+        pre_process = mpu.is_pipeline_first_stage()
+        post_process = mpu.is_pipeline_last_stage()
+
+        # ---- Build GPTModel with MoE ----
+        gpt_config = _make_dsa_moe_config(num_layers=_NUM_GPT_LAYERS, tp=tp, pp=pp)
+        gpt_model = _build_gpt_model(gpt_config, pre_process=pre_process, post_process=post_process)
+        num_local_gpt_layers = len(gpt_model.decoder.layers)
+        gpt_sd = gpt_model.state_dict()
+
+        # ---- Build MambaModel with MoE pattern ----
+        mamba_model = _build_mamba_model(
+            gpt_config, _MOE_MAMBA_PATTERN, pre_process=pre_process, post_process=post_process
+        )
+
+        # ---- Remap GPT weights → Mamba (includes pre_mlp_layernorm.* for MoE layers) ----
+        mamba_sd = _remap_gpt_to_mamba_state_dict(gpt_sd, num_local_gpt_layers)
+        missing, unexpected = mamba_model.load_state_dict(mamba_sd, strict=True)
+        assert not missing, f"Missing keys after weight remap: {missing}"
+        assert not unexpected, f"Unexpected keys after weight remap: {unexpected}"
+
+        # ---- Create identical inputs on all ranks ----
+        torch.manual_seed(99)
+        tokens = torch.randint(0, _VOCAB_SIZE, (_BATCH_SIZE, _SEQ_LEN), device='cuda')
+
+        # ---- Forward pass ----
+        if pp == 1:
+            gpt_logprobs = _forward_logprobs_pp1(gpt_model, tokens)
+            mamba_logprobs = _forward_logprobs_pp1(mamba_model, tokens)
+            torch.testing.assert_close(
+                gpt_logprobs,
+                mamba_logprobs,
+                atol=1e-5,
+                rtol=1e-5,
+                msg=f"MoE logprob mismatch for tp={tp} pp={pp}",
+            )
+        else:
+            gpt_logprobs = _forward_logprobs_pp2(gpt_model, tokens)
+            mamba_logprobs = _forward_logprobs_pp2(mamba_model, tokens)
+            if post_process:
+                torch.testing.assert_close(
+                    gpt_logprobs,
+                    mamba_logprobs,
+                    atol=1e-5,
+                    rtol=1e-5,
+                    msg=f"MoE logprob mismatch for tp={tp} pp={pp}",
+                )
+
+    def test_moe_weight_loading_strict(self, tp: int, pp: int) -> None:
+        """Verify that strict=True weight loading succeeds with MoE keys."""
+        self._skip_if_insufficient_gpus(tp, pp)
+        Utils.initialize_model_parallel(tp, pp)
+        model_parallel_cuda_manual_seed(42)
+
+        pre_process = mpu.is_pipeline_first_stage()
+        post_process = mpu.is_pipeline_last_stage()
+
+        gpt_config = _make_dsa_moe_config(num_layers=_NUM_GPT_LAYERS, tp=tp, pp=pp)
+        gpt_model = _build_gpt_model(gpt_config, pre_process=pre_process, post_process=post_process)
+        mamba_model = _build_mamba_model(
+            gpt_config, _MOE_MAMBA_PATTERN, pre_process=pre_process, post_process=post_process
+        )
+
+        gpt_sd = gpt_model.state_dict()
+        num_local_gpt_layers = len(gpt_model.decoder.layers)
+        mamba_sd = _remap_gpt_to_mamba_state_dict(gpt_sd, num_local_gpt_layers)
+        missing, unexpected = mamba_model.load_state_dict(mamba_sd, strict=True)
+
+        assert not missing, f"Missing keys: {missing}"
+        assert not unexpected, f"Unexpected keys: {unexpected}"
+
+    def test_moe_record_and_compare_golden_values(self, tp: int, pp: int) -> None:
+        """Record GPTModel+MoE logprobs as golden values, then compare MambaModel+MoE."""
+        self._skip_if_insufficient_gpus(tp, pp)
+        if tp != 1 or pp != 1:
+            pytest.skip("Golden-value recording only runs for tp=1, pp=1")
+
+        Utils.initialize_model_parallel(1, 1)
+        model_parallel_cuda_manual_seed(42)
+
+        gpt_config = _make_dsa_moe_config(num_layers=_NUM_GPT_LAYERS, tp=1, pp=1)
+        gpt_model = _build_gpt_model(gpt_config)
+        mamba_model = _build_mamba_model(gpt_config, _MOE_MAMBA_PATTERN)
+
+        gpt_sd = gpt_model.state_dict()
+        mamba_sd = _remap_gpt_to_mamba_state_dict(gpt_sd, len(gpt_model.decoder.layers))
+        mamba_model.load_state_dict(mamba_sd, strict=True)
+
+        torch.manual_seed(99)
+        tokens = torch.randint(0, _VOCAB_SIZE, (_BATCH_SIZE, _SEQ_LEN), device='cuda')
+
+        gpt_logprobs = _forward_logprobs_pp1(gpt_model, tokens)
+        mamba_logprobs = _forward_logprobs_pp1(mamba_model, tokens)
+
+        # Verify MambaModel matches golden values
+        _compare_against_golden_values(mamba_logprobs, gpt_logprobs, abs_tol=1e-3)
diff --git a/tests/unit_tests/ssm/test_mamba_block.py b/tests/unit_tests/ssm/test_mamba_block.py
index 7b743afbfad..65d42ea3c17 100644
--- a/tests/unit_tests/ssm/test_mamba_block.py
+++ b/tests/unit_tests/ssm/test_mamba_block.py
@@ -12,7 +12,10 @@
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.attention import SelfAttention
+from megatron.core.transformer.experimental_attention_variant.dsa import DSAttention
 from megatron.core.transformer.mlp import MLP
+from megatron.core.transformer.multi_latent_attention import MLASelfAttention
+from megatron.core.transformer.transformer_config import MLATransformerConfig
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from tests.unit_tests.test_utilities import Utils
 
@@ -46,6 +49,38 @@ def get_mamba_block(self, layer_pattern):
             pg_collection=self.get_pg_collection(),
         )
 
+    def get_dsa_mamba_block(self, layer_pattern):
+        layer_type_list = validate_segment_layers(layer_pattern)
+        transformer_config = MLATransformerConfig(
+            hidden_size=256,  # The Mamba layer places several constraints on this
+            # Need to specify num_attention_heads and num_layers or TransformerConfig
+            # will generate errors.
+            num_layers=len(layer_type_list),
+            num_attention_heads=16,
+            use_cpu_initialization=True,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            q_lora_rank=64,
+            kv_lora_rank=64,
+            qk_head_dim=64,
+            qk_pos_emb_head_dim=32,
+            v_head_dim=64,
+            rope_type='rope',
+            rotary_base=10000,
+            rotary_percent=1.0,
+            dsa_indexer_n_heads=8,
+            dsa_indexer_head_dim=64,
+            dsa_indexer_topk=32,
+        )
+        modules = mamba_stack_spec.submodules
+        return MambaStack(
+            transformer_config,
+            modules,
+            layer_type_list=layer_type_list,
+            pp_layer_offset=0,
+            pg_collection=self.get_pg_collection(),
+        )
+
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
 
@@ -138,3 +173,20 @@ def test_gdn_gpu_forward(self):
         assert output.shape[1] == micro_batch_size
         assert output.shape[2] == block.config.hidden_size
         assert output.dtype == torch.float32
+
+    def test_dsa_layer_types(self):
+        """D symbol creates a TransformerLayer with MLASelfAttention."""
+        layer_pattern = Symbols.MAMBA + Symbols.DS_ATTENTION + Symbols.MAMBA
+        block = self.get_dsa_mamba_block(layer_pattern)
+        layers = block.layers
+        assert isinstance(layers[0], MambaLayer)
+        assert isinstance(layers[1], TransformerLayer)
+        assert isinstance(layers[1].self_attention, MLASelfAttention)
+        assert isinstance(layers[1].self_attention.core_attention, DSAttention)
+        assert isinstance(layers[2], MambaLayer)
+
+    def test_mixed_attention_and_dsa_layer_types(self):
+        """* and D in the same block fail."""
+        layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.DS_ATTENTION + Symbols.MAMBA
+        with pytest.raises(ValueError):
+            block = self.get_dsa_mamba_block(layer_pattern)
diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
index 440c843bc27..915ba9dd34f 100644
--- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
+++ b/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
+import operator
 from unittest.mock import patch
 
 import pytest
@@ -10,6 +11,7 @@
     get_hybrid_layer_counts,
     get_hybrid_total_layer_count,
     get_hybrid_total_pipeline_segment_count,
+    get_layer_maps_from_layer_type_list,
     parse_hybrid_pattern,
     pattern_from_ratios,
     select_pipeline_segment,
@@ -75,6 +77,7 @@ def test_valid_patterns(self):
             ("", []),
             ("GGG*GGG*", ['G', 'G', 'G', '*', 'G', 'G', 'G', '*']),
             ("GEGEGE*E", ['G', 'E', 'G', 'E', 'G', 'E', '*', 'E']),
+            ("MDMD", ['M', 'D', 'M', 'D']),
         ]
         for pattern, expected in test_cases:
             result = validate_segment_layers(pattern)
@@ -95,6 +98,9 @@ def test_invalid_symbols_cause_failure(self):
             validate_segment_layers("M|M")  # pipe not valid in a segment
         with pytest.raises(ValueError):
             validate_segment_layers("M/M")  # MTP separator not valid in a segment
+        with pytest.raises(ValueError):
+            # Not allowed to have both standard Attention and MLA/DSA
+            validate_segment_layers("MDM*-")
 
 
 @pytest.mark.internal
@@ -155,6 +161,8 @@ def test_main_pattern_only(self):
             ("E", "E"),
             ("GGG*GGG*", "GGG*GGG*"),
             ("GEGEGE*E", "GEGEGE*E"),
+            ("MDMD", "MDMD"),
+            ("DM", "DM"),
         ]
         for pattern, expected_main in test_cases:
             result = parse_hybrid_pattern(pattern)
@@ -277,6 +285,8 @@ def test_complex_patterns(self):
             ("MEME/MM/MM", "MEME", "MM", 2),
             # GDN+MoE main pattern with GDN MTP
             ("GEGEGE*E/GG/GG", "GEGEGE*E", "GG", 2),
+            # DSA in main pattern with MTP
+            ("MDMD/MD/MD", "MDMD", "MD", 2),
         ]
         for pattern, expected_main, expected_mtp, expected_depths in test_cases:
             result = parse_hybrid_pattern(pattern)
@@ -295,47 +305,74 @@ def test_dataclass_equality(self):
 class TestGetHybridLayerCounts:
 
     def test_simple_pattern(self):
-        assert get_hybrid_layer_counts("M*M*") == {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("M*M*") == {'*': 2, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 0}
 
     def test_all_layer_types(self):
-        assert get_hybrid_layer_counts("MG*-E") == {'M': 1, 'G': 1, '*': 1, '-': 1, 'E': 1}
+        # Not allowed to have both standard Attention and MLA/DSA, so we do separate asserts.
+        assert get_hybrid_layer_counts("MG*-E") == {'*': 1, 'D': 0, 'G': 1, 'M': 1, '-': 1, 'E': 1}
+        assert get_hybrid_layer_counts("MGD-E") == {'*': 0, 'D': 1, 'G': 1, 'M': 1, '-': 1, 'E': 1}
 
     def test_with_pipes(self):
         # Pipes should be skipped in counting
-        assert get_hybrid_layer_counts("M*|M*") == {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0}
-        assert get_hybrid_layer_counts("M-M-|M-M*-") == {'M': 4, 'G': 0, '*': 1, '-': 4, 'E': 0}
+        assert get_hybrid_layer_counts("M*|M*") == {'*': 2, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("M-M-|M-M*-") == {
+            '*': 1,
+            'D': 0,
+            'G': 0,
+            'M': 4,
+            '-': 4,
+            'E': 0,
+        }
 
     def test_with_mtp(self):
         # MTP pattern "MM" repeated 2 depths -> 4 extra mamba layers
-        assert get_hybrid_layer_counts("M*M*/MM/MM") == {'M': 6, 'G': 0, '*': 2, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("M*M*/MM/MM") == {
+            '*': 2,
+            'D': 0,
+            'G': 0,
+            'M': 6,
+            '-': 0,
+            'E': 0,
+        }
 
     def test_with_pipes_and_mtp(self):
         # Main: M-M-|M-M*- -> 1 attn, 4 mamba, 4 mlp
         # MTP: MM x 2 depths -> +4 mamba
         assert get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") == {
-            'M': 8,
-            'G': 0,
             '*': 1,
+            'D': 0,
+            'G': 0,
+            'M': 8,
             '-': 4,
             'E': 0,
         }
 
     def test_moe_pattern(self):
-        assert get_hybrid_layer_counts("MEME") == {'M': 2, 'G': 0, '*': 0, '-': 0, 'E': 2}
+        assert get_hybrid_layer_counts("MEME") == {'*': 0, 'D': 0, 'G': 0, 'M': 2, '-': 0, 'E': 2}
 
     def test_mtp_with_attention(self):
         # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP
-        assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == {'M': 7, 'G': 0, '*': 3, '-': 0, 'E': 0}
-
-    def test_empty_pattern(self):
-        assert get_hybrid_layer_counts("") == {'M': 0, 'G': 0, '*': 0, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == {
+            '*': 3,
+            'D': 0,
+            'G': 0,
+            'M': 7,
+            '-': 0,
+            'E': 0,
+        }
 
     def test_gdn_pattern(self):
-        assert get_hybrid_layer_counts("GMGM") == {'M': 2, 'G': 2, '*': 0, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("GMGM") == {'*': 0, 'D': 0, 'G': 2, 'M': 2, '-': 0, 'E': 0}
 
     def test_gdn_hybrid_pattern(self):
         # GDN + Mamba + Attention
-        assert get_hybrid_layer_counts("G*GM*") == {'M': 1, 'G': 2, '*': 2, '-': 0, 'E': 0}
+        assert get_hybrid_layer_counts("G*GM*") == {'*': 2, 'D': 0, 'G': 2, 'M': 1, '-': 0, 'E': 0}
+
+    def test_dsa_pattern(self):
+        assert get_hybrid_layer_counts("DMDM") == {'*': 0, 'D': 2, 'G': 0, 'M': 2, '-': 0, 'E': 0}
+
+    def test_empty_pattern(self):
+        assert get_hybrid_layer_counts("") == {'*': 0, 'D': 0, 'G': 0, 'M': 0, '-': 0, 'E': 0}
 
 
 @pytest.mark.internal
@@ -596,3 +633,56 @@ def test_all_ranks_cover_full_pattern(self):
             assert offset == len(all_layers)
             all_layers.extend(layers)
         assert all_layers == ['M', '*', 'M', '*', 'M', '*']
+
+
+@pytest.mark.internal
+class TestGetLayerMapsFromLayerTypeList:
+    """Tests for get_layer_maps_from_layer_type_list."""
+
+    def test_standard_layer_types(self):
+        """Standard symbols each produce a single-entry map at local index 0."""
+        maps = get_layer_maps_from_layer_type_list(["*", "M", "-", "E"])
+        # We always get all symbols returned, not only those contained in the pattern.
+        assert len(maps) == 6
+        attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter(
+            Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE
+        )(maps)
+        assert attention_map == {0: 0}
+        assert mamba_map == {1: 0}
+        assert mlp_map == {2: 0}
+        assert moe_map == {3: 0}
+
+    def test_dsa(self):
+        """D (DSA) layers are treated as separate layers for KV cache mapping."""
+        maps = get_layer_maps_from_layer_type_list(["D", "M", "D", "M"])
+        attention_map, dsa_map, mamba_map, mlp_map, moe_map = operator.itemgetter(
+            Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE
+        )(maps)
+        assert attention_map == {}
+        assert dsa_map == {0: 0, 2: 1}
+        assert mamba_map == {1: 0, 3: 1}
+        assert mlp_map == {}
+        assert moe_map == {}
+
+    def test_mixed_attention_and_dsa(self):
+        """Both * and D contribute to the different maps with non-consecutive local indices."""
+        maps = get_layer_maps_from_layer_type_list(["*", "D", "M", "-"])
+        attention_map, dsa_map, mamba_map, mlp_map, moe_map = operator.itemgetter(
+            Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE
+        )(maps)
+        assert attention_map == {0: 0}
+        assert dsa_map == {1: 0}
+        assert mamba_map == {2: 0}
+        assert mlp_map == {3: 0}
+        assert moe_map == {}
+
+    def test_all_mamba(self):
+        """All-mamba pattern leaves attention, mlp, and moe maps empty."""
+        maps = get_layer_maps_from_layer_type_list(["M", "M", "M"])
+        attention_map, mamba_map, mlp_map, moe_map = operator.itemgetter(
+            Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE
+        )(maps)
+        assert attention_map == {}
+        assert mamba_map == {0: 0, 1: 1, 2: 2}
+        assert mlp_map == {}
+        assert moe_map == {}
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py
index 45470d4dd6c..6d2426a42d1 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py
@@ -6,6 +6,10 @@
 import torch
 
 import megatron.core.parallel_state as parallel_state
+from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+    get_dsa_module_spec_for_backend,
+    get_experimental_attention_variant_module_spec,
+)
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
@@ -23,6 +27,7 @@
     fused_qk_topk_naive,
     rotate_activation,
 )
+from megatron.core.transformer.multi_latent_attention import MLASelfAttention
 from megatron.core.transformer.transformer_config import MLATransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
@@ -1587,3 +1592,71 @@ def test_dsa_gradient_sync(
                         ), f"Indexer gradient for {name} differs between TP rank 0 and rank {i} after TP sync"
 
         Utils.destroy_model_parallel()
+
+
+@pytest.mark.internal
+class TestDSAModuleSpecDispatch:
+    """Tests for get_dsa_module_spec_for_backend and get_experimental_attention_variant_module_spec."""
+
+    @pytest.fixture(scope='function', autouse=True)
+    def setup_method(self):
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=1, pipeline_model_parallel_size=1
+        )
+        yield
+        Utils.destroy_model_parallel()
+
+    def _make_dsa_config(self, **kwargs):
+        return MLATransformerConfig(
+            num_layers=2,
+            hidden_size=256,
+            num_attention_heads=16,
+            use_cpu_initialization=True,
+            bf16=True,
+            params_dtype=torch.bfloat16,
+            q_lora_rank=64,
+            kv_lora_rank=64,
+            qk_head_dim=64,
+            qk_pos_emb_head_dim=32,
+            v_head_dim=64,
+            rope_type='rope',
+            rotary_base=10000,
+            rotary_percent=1.0,
+            dsa_indexer_n_heads=8,
+            dsa_indexer_head_dim=64,
+            dsa_indexer_topk=32,
+            **kwargs,
+        )
+
+    def test_get_experimental_attention_variant_module_spec_dsa(self):
+        """get_experimental_attention_variant_module_spec dispatches to DSA for variant='dsa'."""
+        config = self._make_dsa_config(experimental_attention_variant="dsa")
+        spec = get_experimental_attention_variant_module_spec(config)
+        assert spec.module == MLASelfAttention
+        assert spec.submodules.core_attention.module == DSAttention
+
+    def test_get_dsa_module_spec_for_backend(self):
+        """get_dsa_module_spec_for_backend returns the correct full spec structure."""
+        from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
+
+        config = self._make_dsa_config()
+        backend = TESpecProvider()
+        spec = get_dsa_module_spec_for_backend(config, backend=backend)
+        assert spec.module == MLASelfAttention
+        assert spec.submodules.core_attention.module == DSAttention
+        assert spec.submodules.core_attention.submodules.indexer.module == DSAIndexer
+        assert spec.params["attn_mask_type"] == AttnMaskType.causal
+
+    def test_get_dsa_module_spec_requires_mla(self):
+        """get_dsa_module_spec_for_backend rejects configs without MLA."""
+        from megatron.core.transformer import TransformerConfig as _TransformerConfig
+
+        config = _TransformerConfig(num_layers=2, hidden_size=256, num_attention_heads=4)
+        with pytest.raises(AssertionError, match="only MLA supports sparse attention"):
+            get_dsa_module_spec_for_backend(config, backend=None)
+
+    def test_get_dsa_module_spec_rejects_qk_l2_norm(self):
+        """get_dsa_module_spec_for_backend rejects configs with qk_l2_norm=True."""
+        config = self._make_dsa_config(qk_l2_norm=True)
+        with pytest.raises(AssertionError, match="qk_l2_norm is not supported"):
+            get_dsa_module_spec_for_backend(config, backend=None)
diff --git a/tools/checkpoint/remap_gpt_dsa_to_mamba.py b/tools/checkpoint/remap_gpt_dsa_to_mamba.py
new file mode 100644
index 00000000000..8a6888d1dc7
--- /dev/null
+++ b/tools/checkpoint/remap_gpt_dsa_to_mamba.py
@@ -0,0 +1,179 @@
+#!/usr/bin/env python3
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+"""Convert a GPTModel DSA checkpoint to a MambaModel-compatible checkpoint.
+
+A GPTModel with ``--experimental-attention-variant dsa`` uses one combined
+TransformerLayer per model layer (attention + MLP).  The equivalent MambaModel
+with pattern ``D-D-...`` stores them as two separate layers:
+
+* Layer 2N   – DSA attention (TransformerLayer: input_layernorm + MLASelfAttention)
+* Layer 2N+1 – MLP          (MLPLayer: fused-norm MLP)
+
+This script loads a GPTModel Distributed Checkpoint (DCP), remaps the state-dict
+keys, and saves a new DCP that can be loaded by MambaModel.
+
+Usage
+-----
+::
+
+    python tools/checkpoint/remap_gpt_dsa_to_mamba.py \\
+        --input  /path/to/gpt_dsa_dcp_checkpoint \\
+        --output /path/to/mamba_dsa_dcp_checkpoint \\
+        --num-gpt-layers 4
+
+Key remapping rules
+-------------------
+* ``decoder.layers.{N}.input_layernorm.*``    →  ``decoder.layers.{2N}.input_layernorm.*``
+* ``decoder.layers.{N}.self_attention.*``     →  ``decoder.layers.{2N}.self_attention.*``
+* ``decoder.layers.{N}.pre_mlp_layernorm.*`` →  ``decoder.layers.{2N+1}.pre_mlp_layernorm.*``
+* ``decoder.layers.{N}.mlp.*``               →  ``decoder.layers.{2N+1}.mlp.*``
+* ``decoder.final_layernorm.*``              →  ``decoder.final_norm.*``
+* All other keys                             →  unchanged
+
+Note: ``pre_mlp_layernorm`` only appears for MoE layers (where it is a real TENorm).
+Dense layers use ``IdentityOp`` for ``pre_mlp_layernorm``, which produces no state dict keys.
+"""
+
+import argparse
+import os
+import re
+import shutil
+from pathlib import Path
+from typing import Dict
+
+
+def _remap_key(key: str, num_gpt_layers: int) -> str:
+    """Return the MambaModel state-dict key corresponding to *key* from GPTModel.
+
+    Args:
+        key: A key from the GPTModel state dict.
+        num_gpt_layers: Total number of GPT decoder layers (across all PP stages).
+
+    Returns:
+        The remapped key for MambaModel.
+
+    Raises:
+        ValueError: If an unexpected sub-key is encountered in a decoder layer.
+    """
+    layer_prefix = "decoder.layers."
+    final_ln_prefix = "decoder.final_layernorm."
+
+    # Final layernorm name differs between TransformerBlock and MambaStack
+    if key.startswith(final_ln_prefix):
+        return "decoder.final_norm." + key[len(final_ln_prefix):]
+
+    if not key.startswith(layer_prefix):
+        return key  # embedding, output_layer, rotary_pos_emb, etc.
+
+    # Parse "decoder.layers.{N}.{rest}"
+    remainder = key[len(layer_prefix):]
+    dot_idx = remainder.index('.')
+    layer_n = int(remainder[:dot_idx])
+    rest = remainder[dot_idx + 1:]
+
+    assert 0 <= layer_n < num_gpt_layers, f"Layer index {layer_n} out of range [0, {num_gpt_layers}) in key '{key}'"
+
+    if rest.startswith("input_layernorm.") or rest.startswith("self_attention."):
+        return f"{layer_prefix}{2 * layer_n}.{rest}"
+    elif rest.startswith("mlp."):
+        return f"{layer_prefix}{2 * layer_n + 1}.{rest}"
+    elif rest.startswith("pre_mlp_layernorm."):
+        # MoE layers have a real TENorm for pre_mlp_layernorm (not fused);
+        # it maps to MoETransformerLayer 2N+1.
+        return f"{layer_prefix}{2 * layer_n + 1}.{rest}"
+    else:
+        raise ValueError(
+            f"Unexpected sub-key '{rest}' in GPT layer {layer_n} (full key='{key}'). "
+            "Expected: input_layernorm.*, self_attention.*, pre_mlp_layernorm.*, mlp.*"
+        )
+
+
+def _remap_state_dict(
+    gpt_sd: Dict, num_gpt_layers: int
+) -> Dict:
+    """Apply key remapping to the full GPTModel state dict."""
+    return {_remap_key(k, num_gpt_layers): v for k, v in gpt_sd.items()}
+
+
+def convert(input_path: Path, output_path: Path, num_gpt_layers: int) -> None:
+    """Load a GPTModel DCP checkpoint, remap keys, and save as MambaModel DCP.
+
+    Args:
+        input_path: Path to the GPTModel DCP checkpoint directory.
+        output_path: Destination directory for the MambaModel DCP checkpoint.
+        num_gpt_layers: Number of GPT decoder layers in the original model.
+    """
+    try:
+        import torch
+        import torch.distributed.checkpoint as dcp
+        from torch.distributed.checkpoint.format_utils import (
+            dcp_to_torch_save,
+            torch_save_to_dcp,
+        )
+    except ImportError as exc:
+        raise SystemExit(
+            "PyTorch distributed checkpoint (torch.distributed.checkpoint) is required. "
+            "Please upgrade to PyTorch >= 2.0."
+        ) from exc
+
+    print(f"Loading GPTModel checkpoint from: {input_path}")
+
+    # --- Load the flat state dict from DCP ---
+    # We use dcp_to_torch_save to materialize the DCP into a regular .pt file,
+    # then remap keys, then convert back to DCP.
+    tmp_flat = output_path.parent / "_tmp_gpt_flat.pt"
+    try:
+        dcp_to_torch_save(str(input_path), str(tmp_flat))
+        gpt_sd = torch.load(tmp_flat, map_location="cpu")
+        print(f"Loaded {len(gpt_sd)} keys from GPTModel checkpoint.")
+
+        # --- Remap keys ---
+        mamba_sd = _remap_state_dict(gpt_sd, num_gpt_layers)
+        print(
+            f"Remapped state dict: {len(gpt_sd)} GPT keys → {len(mamba_sd)} Mamba keys."
+        )
+
+        # --- Save remapped state dict as a new flat .pt then convert to DCP ---
+        tmp_mamba = output_path.parent / "_tmp_mamba_flat.pt"
+        torch.save(mamba_sd, tmp_mamba)
+
+        output_path.mkdir(parents=True, exist_ok=True)
+        torch_save_to_dcp(str(tmp_mamba), str(output_path))
+        print(f"MambaModel DCP checkpoint saved to: {output_path}")
+
+    finally:
+        for tmp in (tmp_flat, output_path.parent / "_tmp_mamba_flat.pt"):
+            if tmp.exists():
+                tmp.unlink()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Convert GPTModel DSA checkpoint to MambaModel-compatible format."
+    )
+    parser.add_argument(
+        "--input", required=True, type=Path,
+        help="Path to the source GPTModel DCP checkpoint directory.",
+    )
+    parser.add_argument(
+        "--output", required=True, type=Path,
+        help="Destination path for the MambaModel DCP checkpoint.",
+    )
+    parser.add_argument(
+        "--num-gpt-layers", required=True, type=int,
+        help="Number of decoder layers in the GPTModel (e.g. 4).",
+    )
+    args = parser.parse_args()
+
+    if not args.input.exists():
+        raise SystemExit(f"Input checkpoint not found: {args.input}")
+    if args.output.exists():
+        print(f"Warning: output path already exists and will be overwritten: {args.output}")
+        shutil.rmtree(args.output)
+
+    convert(args.input, args.output, args.num_gpt_layers)
+    print("Conversion complete.")
+
+
+if __name__ == "__main__":
+    main()