[dev] DeepSeek V3.2 support

gpt_builders.py

@@ -112,6 +112,7 @@ def _get_transformer_layer_spec(use_te, config):
             args.moe_grouped_gemm,
             args.qk_layernorm,
             args.multi_latent_attention,
+            args.sparse_attention_type,
             args.linear_attention_type,
             moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             qk_l2_norm=args.qk_l2_norm,
@@ -123,6 +124,7 @@ def _get_transformer_layer_spec(use_te, config):
             args.moe_grouped_gemm,
             args.qk_layernorm,
             args.multi_latent_attention,
+            args.sparse_attention_type,
             args.linear_attention_type,
             moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             normalization=args.normalization,

megatron/core/models/gpt/gpt_layer_specs.py

@@ -9,6 +9,9 @@
     get_linear_attention_module_spec_for_backend,
 )
 from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend
+from megatron.core.models.gpt.sparse_attention_module_specs import (
+    get_sparse_attention_module_spec_for_backend,
+)
 from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
 from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.identity_op import IdentityOp
@@ -77,6 +80,7 @@ def get_gpt_layer_with_transformer_engine_spec(
     moe_grouped_gemm: Optional[bool] = False,
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
+    sparse_attention_type: Optional[str] = None,
     linear_attention_type: Optional[str] = None,
     fp8: Optional[str] = None,  # pylint: disable=unused-argument
     moe_use_legacy_grouped_gemm: Optional[bool] = False,
@@ -94,6 +98,7 @@ def get_gpt_layer_with_transformer_engine_spec(
         moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False.
         qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
         multi_latent_attention (bool, optional): To use multi-latent attention. Defaults to False.
+        sparse_attention_type (str, optional): The type of sparse attention. Defaults to None.
         linear_attention_type (str, optional): The type of linear attention. Defaults to None.
         fp8 (str, optional): Deprecated. For temporary Nemo compatibility.
         moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP.
@@ -129,6 +134,7 @@ def get_gpt_layer_with_transformer_engine_spec(
     attention = get_attention_module_spec_for_backend(
         backend=backend,
         sharded_state_dict_keys_map=sharded_state_dict_keys_map,
+        sparse_attention_type=sparse_attention_type,
         linear_attention_type=linear_attention_type,
         qk_layernorm=qk_layernorm,
         qk_l2_norm=qk_l2_norm,
@@ -161,6 +167,7 @@ def get_gpt_layer_local_spec(
     moe_grouped_gemm: Optional[bool] = False,
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
+    sparse_attention_type: Optional[str] = None,
     linear_attention_type: Optional[str] = None,
     fp8: Optional[str] = None,  # pylint: disable=unused-argument
     moe_use_legacy_grouped_gemm: Optional[bool] = False,
@@ -176,6 +183,7 @@ def get_gpt_layer_local_spec(
         moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False.
         qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
         multi_latent_attention (bool, optional): To use multi-latent attention. Defaults to False.
+        sparse_attention_type (str, optional): The type of sparse attention. Defaults to None.
         linear_attention_type (str, optional): The type of linear attention. Defaults to None.
         fp8 (str, optional): Deprecated. For temporary Nemo compatibility.
         moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP.
@@ -200,6 +208,9 @@ def get_gpt_layer_local_spec(
             " and will be removed soon. Please update your code accordingly."
         )
 
+    if sparse_attention_type is not None:
+        raise NotImplementedError("Sparse attention is not supported with local spec yet.")
+
     if linear_attention_type is not None:
         raise NotImplementedError("Linear attention is not supported with local spec yet.")
 
@@ -208,6 +219,7 @@ def get_gpt_layer_local_spec(
     attention = get_attention_module_spec_for_backend(
         backend=backend,
         sharded_state_dict_keys_map=sharded_state_dict_keys_map,
+        sparse_attention_type=sparse_attention_type,
         linear_attention_type=linear_attention_type,
         qk_layernorm=qk_layernorm,
         qk_l2_norm=qk_l2_norm,
@@ -272,6 +284,7 @@ def get_transformer_layer_spec_for_backend(
 def get_attention_module_spec_for_backend(
     backend: BackendSpecProvider,
     sharded_state_dict_keys_map: dict,
+    sparse_attention_type: Optional[str] = None,
     linear_attention_type: Optional[str] = None,
     qk_layernorm: Optional[bool] = False,
     qk_l2_norm: Optional[bool] = False,
@@ -292,6 +305,16 @@ def get_attention_module_spec_for_backend(
     rms_norm = normalization == "RMSNorm"
     qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True)
 
+    if sparse_attention_type is not None:
+        assert multi_latent_attention, "Currently only MLA supports sparse attention."
+        core_attention = get_sparse_attention_module_spec_for_backend(
+            backend=backend,
+            sparse_attention_type=sparse_attention_type,
+            normalization=normalization,
+        )
+    else:
+        core_attention = backend.core_attention()
+
     if multi_latent_attention:
         assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA."
         linear_q_down_proj = (
@@ -328,7 +351,7 @@ def get_attention_module_spec_for_backend(
                 linear_q_up_proj=linear_q_up_proj,
                 linear_kv_down_proj=linear_kv_down_proj,
                 linear_kv_up_proj=linear_kv_up_proj,
-                core_attention=backend.core_attention(),
+                core_attention=core_attention,
                 linear_proj=backend.row_parallel_linear(),
                 q_layernorm=qk_norm,
                 kv_layernorm=qk_norm,
@@ -352,7 +375,7 @@ def get_attention_module_spec_for_backend(
             params={"attn_mask_type": AttnMaskType.causal},
             submodules=SelfAttentionSubmodules(
                 linear_qkv=linear_qkv,
-                core_attention=backend.core_attention(),
+                core_attention=core_attention,
                 linear_proj=backend.row_parallel_linear(),
                 q_layernorm=qk_norm,
                 k_layernorm=qk_norm,
@@ -522,15 +545,18 @@ def get_gpt_decoder_block_spec(
                     continue
                 linear_attention_type = config.linear_attention_type
                 multi_latent_attention = None
+                sparse_attention_type = None
             else:
                 linear_attention_type = None
                 multi_latent_attention = config.multi_latent_attention
+                sparse_attention_type = config.sparse_attention_type
 
             layer_spec_key = f"{mlp_type}_{attention_type}"
             layer_spec_dict[layer_spec_key] = get_layer_spec_fn(
                 num_experts=num_experts,
                 moe_grouped_gemm=moe_grouped_gemm,
                 multi_latent_attention=multi_latent_attention,
+                sparse_attention_type=sparse_attention_type,
                 linear_attention_type=linear_attention_type,
                 **get_layer_spec_kwargs,
             )

megatron/core/models/gpt/sparse_attention_module_specs.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Optional
+
+from megatron.core.models.backends import BackendSpecProvider
+from megatron.core.transformer.sparse_attention import (
+    Indexer,
+    IndexerSubmodules,
+    SparseAttention,
+    SparseAttentionSubmodules,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+
+
+def get_indexer_spec_for_backend(
+    backend: BackendSpecProvider, normalization: Optional[str] = None
+) -> ModuleSpec:
+    """Helper function to get Indexer module spec for a given backend.
+
+    Args:
+        backend: Backend specification provider (TE or Local).
+        normalization: Normalization type ("RMSNorm" or None for LayerNorm).
+
+    Returns:
+        ModuleSpec for Indexer with appropriate submodules.
+    """
+    rms_norm = normalization == "RMSNorm"
+    return ModuleSpec(
+        module=Indexer,
+        submodules=IndexerSubmodules(
+            linear_wq_b=backend.linear(),
+            linear_wk=backend.linear(),
+            k_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=True),
+            linear_weights_proj=backend.linear(),
+        ),
+    )
+
+
+def get_sparse_attention_module_spec_for_backend(
+    backend: BackendSpecProvider, sparse_attention_type: str, normalization: Optional[str] = None
+) -> ModuleSpec:
+    """Helper function to get module spec for Sparse Attention.
+
+    Args:
+        backend: Backend specification provider (TE or Local).
+        sparse_attention_type: Type of sparse attention.
+        normalization: Normalization type ("RMSNorm" or None for LayerNorm).
+
+    Returns:
+        ModuleSpec for the sparse attention implementation with appropriate submodules.
+    """
+    if sparse_attention_type == "dsa":
+        # Because TransformerEngine does not support sparse attention yet, we use local
+        # implementation whether the backend is TransformerEngine or not.
+        return ModuleSpec(
+            module=SparseAttention,
+            submodules=SparseAttentionSubmodules(
+                indexer=get_indexer_spec_for_backend(backend, normalization=normalization),
+            ),
+        )
+    else:
+        raise ValueError(f"Invalid sparse attention type: {sparse_attention_type}")

megatron/core/transformer/attention.py

@@ -172,6 +172,7 @@ def __init__(
         self.key_hidden_size = self.hidden_size_per_attention_head
         self.val_hidden_size = self.hidden_size_per_attention_head
 
+        # TODO: This is built twice when using MLA, should be refactored.
         self.core_attention = build_module(
             submodules.core_attention,
             config=self.config,
@@ -188,6 +189,7 @@ def __init__(
             and "core_attn" in self.config.recompute_modules
         )
 
+        # TODO: This is built twice when using MLA, should be refactored.
         # Output.
         self.linear_proj = build_module(
             submodules.linear_proj,

megatron/core/transformer/multi_latent_attention.py

@@ -233,13 +233,22 @@ def forward(
         # Get the query, key and value tensors based on the type of attention -
         # self or cross attn.
         # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128]
-        query, key, value = self.get_query_key_value_tensors(
-            hidden_states,
-            key_value_states,
-            position_ids,
-            packed_seq_params,
-            inference_context=inference_context,
-        )
+        if self.config.sparse_attention_type is None:
+            query, key, value = self.get_query_key_value_tensors(
+                hidden_states,
+                key_value_states,
+                position_ids,
+                packed_seq_params,
+                inference_context=inference_context,
+            )
+        else:
+            query, key, value, q_compressed, _ = self.get_query_key_value_and_compressed_tensors(
+                hidden_states,
+                key_value_states,
+                position_ids,
+                packed_seq_params,
+                inference_context=inference_context,
+            )
 
         # ===================================================
         # Adjust key, value for inference
@@ -267,14 +276,28 @@ def forward(
             )
         else:
             if inference_context is None or inference_context.is_static_batching():
-                core_attn_out = self.core_attention(
-                    query,
-                    key,
-                    value,
-                    attention_mask,
-                    packed_seq_params=packed_seq_params,
-                    attn_mask_type=attn_mask_type,
-                )
+                if self.config.sparse_attention_type is None:
+                    core_attn_out = self.core_attention(
+                        query,
+                        key,
+                        value,
+                        attention_mask,
+                        packed_seq_params=packed_seq_params,
+                        attn_mask_type=attn_mask_type,
+                    )
+                else:
+                    # For sparse attention, use a specialized forward.
+                    core_attn_out = self.core_attention(
+                        query,
+                        key,
+                        value,
+                        x=hidden_states,
+                        qr=q_compressed,
+                        attention_mask=attention_mask,
+                        attn_mask_type=attn_mask_type,
+                        attention_bias=None,
+                        packed_seq_params=packed_seq_params,
+                    )
             elif self.cache_mla_latents:
                 # Dynamic batching attention kernel.
                 q, k, v = (query, key, value)
@@ -460,7 +483,7 @@ def __init__(
             eps=self.config.layernorm_epsilon,
         )
 
-    def get_query_key_value_tensors(
+    def get_query_key_value_and_compressed_tensors(
         self,
         hidden_states,
         key_value_states=None,
@@ -578,6 +601,16 @@ def get_query_key_value_tensors(
             kv_compressed = kv_compressed.squeeze(1)
             k_pos_emb = k_pos_emb.squeeze(1)
 
+        # =========================================
+        # Apply norm
+        # =========================================
+
+        if self.config.q_lora_rank is not None:
+            # q_compressed: [num_tokens, q_lora_rank]
+            q_compressed = self.q_layernorm(q_compressed)
+
+        kv_compressed = self.kv_layernorm(kv_compressed)
+
         # =========================================
         # QKV up projection and RoPE apply
         # =========================================
@@ -588,7 +621,6 @@ def qkv_up_proj_and_rope_apply_for_cached_latent_kv(
             if self.config.q_lora_rank is not None:
                 # q_compressed: [num_tokens, q_lora_rank]
                 # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)]
-                q_compressed = self.q_layernorm(q_compressed)
                 q, _ = self.linear_q_up_proj(q_compressed)
             else:
                 # q_compressed: [num_tokens, hidden_size]
@@ -598,8 +630,6 @@ def qkv_up_proj_and_rope_apply_for_cached_latent_kv(
             # q: [num_tokens, n, q_head_dim]
             q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim)
 
-            kv_compressed = self.kv_layernorm(kv_compressed)
-
             # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim]
             k_pos_emb = torch.unsqueeze(k_pos_emb, -2)
 
@@ -663,7 +693,6 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
             if self.config.q_lora_rank is not None:
                 # q_compressed: [num_tokens, q_lora_rank]
                 # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)]
-                q_compressed = self.q_layernorm(q_compressed)
                 q, _ = self.linear_q_up_proj(q_compressed)
             else:
                 # q_compressed: [num_tokens, hidden_size]
@@ -673,8 +702,6 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
             # q: [num_tokens, n, q_head_dim]
             q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim)
 
-            kv_compressed = self.kv_layernorm(kv_compressed)
-
             # kv: [num_tokens, n * (qk_head_dim + v_head_dim)]
             kv, _ = self.linear_kv_up_proj(kv_compressed)
 
@@ -798,7 +825,11 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po
                     q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb
                 )
 
-        return query, key, value
+        return query, key, value, q_compressed, kv_compressed
+
+    def get_query_key_value_tensors(self, *args, **kwargs):
+        # Only return query, key and value.
+        return get_query_key_value_and_compressed_tensors(self, *args, **kwargs)[:3]
 
     def uncompress_kv_from_cache(self, kv_cached):
         """