Port DeepSeek Sparse Attention to MambaModel

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

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}"

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

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
             )

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

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

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

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: