Port DeepSeek Sparse Attention to MambaModel

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,
 )
@@ -19,7 +20,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,
@@ -96,6 +108,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
@@ -156,6 +203,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/ssm/mamba_block.py

@@ -41,6 +41,7 @@ class MambaStackSubmodules:
 
     mamba_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
@@ -160,6 +161,15 @@ def __init__(
                         pg_collection=pg_collection,
                         is_mtp_layer=is_mtp_layer,
                     )
+                elif layer_type == LayerSymbols.DSA_ATTENTION:
+                    # DSA attention layers apply their own pp_layer_offset
+                    layer = build_module(
+                        submodules.dsa_layer,
+                        config=self.config,
+                        layer_number=i + 1,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
+                    )
                 elif layer_type == LayerSymbols.MLP:
                     # MLP layers apply their own pp_layer_offset
                     layer = build_module(

megatron/core/ssm/mamba_hybrid_layer_allocation.py

@@ -28,10 +28,11 @@ class Symbols:
 
     MAMBA = "M"
     ATTENTION = "*"
+    DSA_ATTENTION = "S"
     MLP = "-"
     MOE = 'E'
     MTP_SEPARATOR = "/"
-    VALID = {MAMBA, ATTENTION, MLP, MOE}
+    VALID = {MAMBA, ATTENTION, DSA_ATTENTION, MLP, MOE}
 
 
 @dataclass
@@ -264,6 +265,7 @@ def allocate_layers(
             logging.INFO,
             f"Hybrid allocation ({Symbols.MAMBA} is mamba, "
             f"{Symbols.ATTENTION} is attention, "
+            f"{Symbols.DSA_ATTENTION} is dsa attention, "
             f"{Symbols.MLP} is mlp):",
         )
         maybe_log_single_rank(logger, logging.INFO, allocation_string)
@@ -303,7 +305,9 @@ def get_layer_maps_from_layer_type_list(
     layer_types = [Symbols.ATTENTION, Symbols.MAMBA, Symbols.MLP, Symbols.MOE]
     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]
+        # DSA attention layers are treated as attention for KV cache mapping.
+        effective_type = Symbols.ATTENTION if layer_type == Symbols.DSA_ATTENTION else layer_type
+        layer_map = layer_maps[effective_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]

megatron/core/ssm/mamba_mixer.py

@@ -34,6 +34,7 @@
     make_sharded_tensors_for_checkpoint,
     sharded_state_dict_default,
 )
+from megatron.core.fp8_utils import get_fp8_align_size
 from megatron.core.utils import (
     deprecate_inference_params,
     is_causal_conv1d_min_version,
@@ -207,9 +208,10 @@ def __init__(
             self.nheads = self.d_inner // self.headdim
 
         if self.config.fp8:
-            assert (2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads) % 16 == 0, (
+            fp8_align_size = get_fp8_align_size(self.config.fp8_recipe)
+            assert (2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads) % fp8_align_size == 0, (
                 "For FP8, the innermost dimension of the Mamba layer "
-                "input projection output tensor must be a multiple of 16."
+                f"input projection output tensor must be a multiple of {fp8_align_size}."
             )
 
         tp_size = self.pg_collection.tp.size()

tests/functional_tests/test_cases/hybrid/hybrid_dsa_mamba_logitsmatch_tp1_pp1/golden_values_dev_dgx_h100.json

@@ -0,0 +1 @@
+{"0": {"input_prompt": "The quick brown fox jumps over the lazy dog.", "generated_text": "", "generated_tokens": [], "logprobs": []}}

tests/functional_tests/test_cases/hybrid/hybrid_dsa_mamba_logitsmatch_tp1_pp1/model_config.yaml

@@ -0,0 +1,49 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: ":4096:8"
+TEST_TYPE: frozen-start
+MODE: inference
+MODEL_ARGS:
+  --load: ${CHECKPOINT_LOAD_PATH}/model/deepseek_v3_proxy/dcp/checkpoint
+  --tokenizer-type: TikTokenizer
+  --tiktoken-pattern: v2
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --use-mcore-models: true
+  --is-hybrid-model: true
+  --model-provider: mamba
+  --disable-bias-linear: true
+  --position-embedding-type: rope
+  --multi-latent-attention: true
+  --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
+  --num-layers: 8
+  --hidden-size: 256
+  --num-attention-heads: 16
+  --hybrid-override-pattern: "S-S-S-S-"
+  --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec
+  --normalization: RMSNorm
+  --swiglu: true
+  --bf16: true
+  --attention-backend: flash
+  --deterministic-mode: true
+  --temperature: 1.0
+  --top_k: 1
+  --return-log-probs: true
+  --num-tokens-to-generate: 30
+  --inference-max-seq-length: 256
+  --output-path: ${INFERENCE_OUTPUT_PATH}
+  --prompts: "The quick brown fox jumps over the lazy dog."
+  --incoming-requests-per-sec: -1
+METRICS:
+  - "generated_tokens"
+  - "logprobs"

tests/functional_tests/test_cases/hybrid/hybrid_dsa_mamba_logitsmatch_tp2_pp1/golden_values_dev_dgx_h100.json

@@ -0,0 +1 @@
+{"0": {"input_prompt": "The quick brown fox jumps over the lazy dog.", "generated_text": "", "generated_tokens": [], "logprobs": []}}

tests/functional_tests/test_cases/hybrid/hybrid_dsa_mamba_logitsmatch_tp2_pp1/model_config.yaml

@@ -0,0 +1,49 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: ":4096:8"
+TEST_TYPE: frozen-start
+MODE: inference
+MODEL_ARGS:
+  --load: ${CHECKPOINT_LOAD_PATH}/model/deepseek_v3_proxy/dcp/checkpoint
+  --tokenizer-type: TikTokenizer
+  --tiktoken-pattern: v2
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 2
+  --pipeline-model-parallel-size: 1
+  --use-mcore-models: true
+  --is-hybrid-model: true
+  --model-provider: mamba
+  --disable-bias-linear: true
+  --position-embedding-type: rope
+  --multi-latent-attention: true
+  --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
+  --num-layers: 8
+  --hidden-size: 256
+  --num-attention-heads: 16
+  --hybrid-override-pattern: "S-S-S-S-"
+  --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec
+  --normalization: RMSNorm
+  --swiglu: true
+  --bf16: true
+  --attention-backend: flash
+  --deterministic-mode: true
+  --temperature: 1.0
+  --top_k: 1
+  --return-log-probs: true
+  --num-tokens-to-generate: 30
+  --inference-max-seq-length: 256
+  --output-path: ${INFERENCE_OUTPUT_PATH}
+  --prompts: "The quick brown fox jumps over the lazy dog."
+  --incoming-requests-per-sec: -1
+METRICS:
+  - "generated_tokens"
+  - "logprobs"

tests/functional_tests/test_cases/hybrid/hybrid_dsa_moe_mamba_logitsmatch_tp1_pp1/golden_values_dev_dgx_h100.json

@@ -0,0 +1 @@
+{"0": {"input_prompt": "The quick brown fox jumps over the lazy dog.", "generated_text": "", "generated_tokens": [], "logprobs": []}}

tests/functional_tests/test_cases/hybrid/hybrid_dsa_moe_mamba_logitsmatch_tp1_pp1/model_config.yaml

@@ -0,0 +1,59 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: ":4096:8"
+TEST_TYPE: frozen-start
+MODE: inference
+MODEL_ARGS:
+  --load: ${CHECKPOINT_LOAD_PATH}/model/deepseek_v3_proxy/dcp/checkpoint
+  --tokenizer-type: TikTokenizer
+  --tiktoken-pattern: v2
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --use-mcore-models: true
+  --is-hybrid-model: true
+  --model-provider: mamba
+  --disable-bias-linear: true
+  --position-embedding-type: rope
+  --multi-latent-attention: true
+  --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
+  --num-layers: 8
+  --hidden-size: 256
+  --num-attention-heads: 16
+  --hybrid-override-pattern: "S-S-SESE"
+  --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec
+  --normalization: RMSNorm
+  --swiglu: true
+  --bf16: true
+  --attention-backend: flash
+  --deterministic-mode: true
+  --temperature: 1.0
+  --top_k: 1
+  --return-log-probs: true
+  --num-tokens-to-generate: 30
+  --inference-max-seq-length: 256
+  --output-path: ${INFERENCE_OUTPUT_PATH}
+  --prompts: "The quick brown fox jumps over the lazy dog."
+  --incoming-requests-per-sec: -1
+  # MoE args
+  --num-experts: 4
+  --moe-layer-freq: ([0]*2+[1]*2)
+  --moe-ffn-hidden-size: 512
+  --moe-shared-expert-intermediate-size: 512
+  --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
+METRICS:
+  - "generated_tokens"
+  - "logprobs"

tests/functional_tests/test_cases/hybrid/hybrid_dsa_moe_mamba_logitsmatch_tp2_pp1/golden_values_dev_dgx_h100.json

@@ -0,0 +1 @@
+{"0": {"input_prompt": "The quick brown fox jumps over the lazy dog.", "generated_text": "", "generated_tokens": [], "logprobs": []}}