[train]: Add DeepSeek-V4 TFLOPs. (#1230)

### PR Category
<!-- One of [ Train | Inference | Compress | Serve | RL | Core |
Hardware | CICD | Tools | Others ] -->
[Train]
### PR Types
<!-- One of [ User Experience | New Features | Bug Fixes | Improvements
| Performance | Breaking Change| Deprecations | Test Case | Docs |
Others ] -->
[Improvements]
Add DeepSeek-V4 TFLOPs. Chane IdentityOp of Engram to None when not
use_engram.
### PR Description
<!-- Describe what you’ve done -->

Author: LiJunscs

flagscale/models/megatron/deepseek_v4/deepseek_builder.py

@@ -108,7 +108,7 @@ def get_deepseek_layer_spec(
     if build_engram:
         engram_module = EngramModule 
     else:
-        engram_module = IdentityOp
+        engram_module = None
     submodules = DeepSeekTransformerLayerSubmodules(
         input_layernorm=input_layernorm,
         self_attention=hybrid_attn_spec,

flagscale/models/megatron/deepseek_v4/deepseek_transformer_layer.py

@@ -41,12 +41,15 @@ def __init__(self, config, submodules, *args, **kwargs):
                 "DeepSeekTransformerLayer now requires config.enable_hyper_connections=True."
             )
         super().__init__(config=config, submodules=submodules, *args, **kwargs)
-
-        self.engram = build_module(
-            submodules.engram,
-            engram_cfg=self.config,
-            layer_id=self.layer_number - 1,
-        )
+        if self.config.use_engram:
+            # If not use_engram, the submodules.engram is None
+            self.engram = build_module(
+                submodules.engram,
+                engram_cfg=self.config,
+                layer_id=self.layer_number - 1,
+            )
+        else:
+            self.engram = None
         self._deepseek_engram_hash_input_ids = None
         self._mhc_recompute_manager = None
         if self.config.engram_layer_ids is not None and self.layer_number - 1 in self.config.engram_layer_ids:
@@ -104,7 +107,7 @@ def _forward_attention(
         inference_params: Optional[Any] = None,
     ):
         """Apply DeepSeek engram before the parent attention path."""
-        if not isinstance(self.engram, IdentityOp):
+        if self.engram is not None:
             nvtx_range_push(suffix="engram")
             hidden_states = self.engram(hidden_states, self._deepseek_engram_hash_input_ids)
             nvtx_range_pop(suffix="engram")

flagscale/train/megatron/training/arguments.py

@@ -3271,6 +3271,13 @@ def _add_experimental_attention_variant_args(parser):
                             "([128]+[4]*23)": 1 HCA layer followed by 23 CSA layers
                             "([128]*3+[4]*2)*2": Three HCA layers followed by two CSA layers, repeated twice.
                         """)
+    group.add_argument(
+        '--no-dsa-kernel-fusion',
+        action='store_false',
+        help='Disable fused DSA sparse-attention kernels (FlashMLA + cuDNN DSA) '
+        'and fall back to unfused PyTorch implementations.',
+        dest='apply_dsa_kernel_fusion',
+    )
     return parser
 
 def _add_heterogeneous_args(parser):

flagscale/train/megatron/training/training.py

@@ -558,42 +558,68 @@ def transformer_flops():
             https://arxiv.org/abs/2305.10403
             https://arxiv.org/abs/2205.05198
             '''
-            ## MLA
-            if args.q_lora_rank is None:
-                q_term = (
-                    args.hidden_size
-                    * args.num_attention_heads
-                    * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                )
-            else:
+            if args.experimental_attention_variant == "dsv4_hybrid":
+                ## DSv4 hybrid MLA projections (per layer, per token).
+                ## In dsv4_hybrid mode, qk_head_dim + qk_pos_emb_head_dim == v_head_dim
+                ## (qk_head_dim is derived as v_head_dim - qk_pos_emb_head_dim), and the
+                ## joint KV is produced by a single hidden -> v_head_dim projection.
+                ## Full core attention is replaced by sparse attention and is accounted
+                ## for in the dsv4_hybrid branch below.
                 q_term = args.q_lora_rank * (
                     args.hidden_size
-                    + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                    + 1
+                    + args.num_attention_heads * args.v_head_dim
+                    + 1  # q norm
                 )
-            standard_self_attn_term = (
-                forward_backward_expansion_factor
-                * fma_expansion_factor
-                * (
-                    ## q lora + rope + q norm
-                    q_term
-                    ## kv lora + rope + kv norm
-                    + args.kv_lora_rank
-                    * (
+                kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim  # kv proj + kv norm
+                ## Grouped low-rank output projection:
+                ##  wo_a:        (n_head * v_head_dim) -> (o_groups * o_lora_rank)
+                ##  linear_proj: (o_groups * o_lora_rank) -> hidden
+                o_term = (
+                    args.num_attention_heads * args.v_head_dim * args.o_lora_rank
+                    + args.o_groups * args.o_lora_rank * args.hidden_size
+                )
+                standard_self_attn_term = (
+                    forward_backward_expansion_factor
+                    * fma_expansion_factor
+                    * (q_term + kv_term + o_term)
+                )
+            else:
+                ## MLA
+                if args.q_lora_rank is None:
+                    q_term = (
+                        args.hidden_size
+                        * args.num_attention_heads
+                        * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                    )
+                else:
+                    q_term = args.q_lora_rank * (
                         args.hidden_size
-                        + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim)
+                        + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
                         + 1
                     )
-                    + args.hidden_size * args.qk_pos_emb_head_dim
-                    ## o proj
-                    + (args.num_attention_heads * args.v_head_dim) * args.hidden_size
-                    ## core attn
-                    + args.seq_length
-                    * (args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim))
-                    / 2  # causal mask (only half of the mask is non-zero)
-                    + args.seq_length * args.num_attention_heads * args.v_head_dim / 2
+                standard_self_attn_term = (
+                    forward_backward_expansion_factor
+                    * fma_expansion_factor
+                    * (
+                        ## q lora + rope + q norm
+                        q_term
+                        ## kv lora + rope + kv norm
+                        + args.kv_lora_rank
+                        * (
+                            args.hidden_size
+                            + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim)
+                            + 1
+                        )
+                        + args.hidden_size * args.qk_pos_emb_head_dim
+                        ## o proj
+                        + (args.num_attention_heads * args.v_head_dim) * args.hidden_size
+                        ## core attn
+                        + args.seq_length
+                        * (args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim))
+                        / 2  # causal mask (only half of the mask is non-zero)
+                        + args.seq_length * args.num_attention_heads * args.v_head_dim / 2
+                    )
                 )
-            )
 
         else:
             ## MHA or GQA