[Dev] restore DSv4 tflops calc in training and fix the packed seq case

megatron/training/training.py

@@ -652,50 +652,81 @@ 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
                 )
-            # Token-linear part of MLA self-attention (lora projs, kv proj, RoPE, output proj).
-            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)
+                )
+                # Sparse attention replaces full core attention; its cost is captured
+                # in dsv4_hybrid_extra_term below.
+                standard_self_attn_core_term = 0
+            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
+                # Token-linear part of MLA self-attention (lora projs, kv proj, RoPE, output proj).
+                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-attention (L^2) part: ``QK^T`` and ``(softmax(QK^T)) V``. The
-            # ``/2`` accounts for the causal mask and the ``*2`` cancels it via FMA.
-            standard_self_attn_core_term = (
-                forward_backward_expansion_factor
-                * fma_expansion_factor
-                * (
-                    args.num_attention_heads
-                    * (args.qk_head_dim + args.qk_pos_emb_head_dim)
-                    / 2
-                    + args.num_attention_heads * args.v_head_dim / 2
+                # Core-attention (L^2) part: ``QK^T`` and ``(softmax(QK^T)) V``. The
+                # ``/2`` accounts for the causal mask and the ``*2`` cancels it via FMA.
+                standard_self_attn_core_term = (
+                    forward_backward_expansion_factor
+                    * fma_expansion_factor
+                    * (
+                        args.num_attention_heads
+                        * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                        / 2
+                        + args.num_attention_heads * args.v_head_dim / 2
+                    )
                 )
-            )
 
         else:
             ## MHA or GQA
@@ -730,6 +761,8 @@ def transformer_flops():
                 * 2  # QK^T and (QK^T)V
             )
 
+        dsv4_hybrid_extra_term = 0
+        dsv4_hybrid_extra_core_term = 0
         if is_linear_attention_variant(args.experimental_attention_variant):
             # Calculate number of dense and MoE Transformer MLPs.
             if isinstance(args.linear_attention_freq, int):
@@ -792,6 +825,108 @@ def transformer_flops():
                     "Invalid experimental_attention_variant: "
                     f"{args.experimental_attention_variant}"
                 )
+        elif args.experimental_attention_variant == "dsv4_hybrid":
+            # DSv4 hybrid: full core attention is replaced by sparse attention (CSA),
+            # and selected layers additionally run a learned indexer (DSA).
+            # The MLA-style projection cost per layer is captured in
+            # ``standard_self_attn_term`` above; here we add the extra per-layer FLOPs
+            # for sparse attention, the main compressor, and the indexer.
+            num_linear_attention_layers = 0
+            linear_self_attn_term = 0
+            num_standard_attention_layers = num_layers
+
+            compress_ratios = args.csa_compress_ratios
+            assert compress_ratios is not None, (
+                "csa_compress_ratios must be set for dsv4_hybrid"
+            )
+            assert len(compress_ratios) == num_layers, (
+                f"Invalid length of csa_compress_ratios: {len(compress_ratios)}, "
+                f"expected num_layers + mtp_num_layers ({num_layers})."
+            )
+            # ratio == 0: window-only (no compressor, no indexer)
+            # ratio == 4: window + learned-topk over compressed KV (compressor + indexer)
+            # ratio == 128: window + all compressed KV (compressor only)
+            n_layers_r0 = sum(1 for r in compress_ratios if r == 0)
+            n_layers_r4 = sum(1 for r in compress_ratios if r == 4)
+            n_layers_r128 = sum(1 for r in compress_ratios if r == 128)
+
+            n_head = args.num_attention_heads
+            v_head_dim = args.v_head_dim
+            window = args.csa_window_size
+            seq_len = args.seq_length
+
+            # ---- Sparse attention (replaces full core attention) ----
+            # Split into token-linear parts (window attention, constant per token)
+            # and L^2 parts (compressed-KV attention, scales with sequence length)
+            # so THD packed sequences get correct seqlen_squared_sum_in_batch scaling.
+
+            # r=0: window-only, fixed per-token cost.
+            sparse_attn_r0 = n_layers_r0 * n_head * window * v_head_dim * 2
+
+            # r=128: window (token-linear) + all compressed KV (L^2).
+            # Compressed positions per token ≈ L/(128*2) (causal /2), x2 for
+            # QK^T + softmax@V → L^2 coefficient: n_head * v_head_dim / 128.
+            sparse_attn_r128_window = n_layers_r128 * n_head * window * v_head_dim * 2
+            sparse_attn_r128_core = n_layers_r128 * n_head * v_head_dim / 128
+
+            # ---- Main compressor (ratio > 0 layers) ----
+            # Two projections per layer (wkv + wgate): hidden -> coff * v_head_dim.
+            # ratio == 4: coff = 2 (overlapping windows)
+            # ratio == 128: coff = 1 (non-overlapping)
+            main_compressor_term = (
+                n_layers_r4 * args.hidden_size * (2 * v_head_dim) * 2
+                + n_layers_r128 * args.hidden_size * (1 * v_head_dim) * 2
+            )
+
+            # ---- r=4 layers: sparse attention + indexer ----
+            if n_layers_r4 > 0:
+                assert args.dsa_indexer_n_heads is not None, (
+                    "dsa_indexer_n_heads must be set for dsv4_hybrid with ratio==4 layers."
+                )
+                assert args.dsa_indexer_head_dim is not None, (
+                    "dsa_indexer_head_dim must be set for dsv4_hybrid with ratio==4 layers."
+                )
+                assert args.dsa_indexer_topk is not None, (
+                    "dsa_indexer_topk must be set for dsv4_hybrid with ratio==4 layers."
+                )
+                idx_n_heads = args.dsa_indexer_n_heads
+                idx_head_dim = args.dsa_indexer_head_dim
+                idx_topk = args.dsa_indexer_topk
+
+                effective_topk_4 = min(idx_topk, seq_len // 4)
+                avg_comp_4 = effective_topk_4 * (1 - effective_topk_4 * 4 / (2 * seq_len))
+                sparse_attn_r4 = (
+                    n_layers_r4 * n_head * (window + avg_comp_4) * v_head_dim * 2
+                )
+
+                # Indexer token-linear: compressor (coff=2, wkv + wgate), Q proj,
+                # weights proj.
+                indexer_token_term = (
+                    n_layers_r4 * args.hidden_size * (2 * idx_head_dim) * 2
+                    + n_layers_r4 * args.q_lora_rank * idx_n_heads * idx_head_dim
+                    + n_layers_r4 * args.hidden_size * idx_n_heads
+                )
+                # Indexer L^2: scoring each query against ~L/4 compressed positions.
+                indexer_scoring_core = n_layers_r4 * idx_n_heads * idx_head_dim / 4
+            else:
+                sparse_attn_r4 = 0
+                indexer_token_term = 0
+                indexer_scoring_core = 0
+
+            sparse_attn_token_term = (
+                sparse_attn_r0 + sparse_attn_r4 + sparse_attn_r128_window
+            )
+
+            dsv4_hybrid_extra_term = (
+                forward_backward_expansion_factor
+                * fma_expansion_factor
+                * (sparse_attn_token_term + main_compressor_term + indexer_token_term)
+            )
+            dsv4_hybrid_extra_core_term = (
+                forward_backward_expansion_factor
+                * fma_expansion_factor
+                * (sparse_attn_r128_core + indexer_scoring_core)
+            )
         else:
             num_linear_attention_layers = 0
             linear_self_attn_term = 0
@@ -802,9 +937,14 @@ def transformer_flops():
         self_attn_term = (
             linear_self_attn_term * num_linear_attention_layers
             + standard_self_attn_term * num_standard_attention_layers
+            + dsv4_hybrid_extra_term
+        )
+        # Core attention (L^2) FLOPs. Standard attention has a uniform per-layer
+        # coefficient; DSv4 sparse attention varies by layer type and is pre-summed.
+        self_attn_core_term = (
+            standard_self_attn_core_term * num_standard_attention_layers
+            + dsv4_hybrid_extra_core_term
         )
-        # Core attention (L^2) FLOPs per standard-attention layer.
-        self_attn_core_term = standard_self_attn_core_term * num_standard_attention_layers
 
         # Token-linear FLOPs scale with the real (unpadded) token count.
         # For BSHD this falls back to ``batch_size * seq_length`` (no padding).

tests/unit_tests/test_num_floating_point_operations.py

@@ -644,3 +644,133 @@ def test_dedup_across_topology(self, tp, cp, pp):
             f"topology tp={tp} cp={cp} pp={pp} dp={dp_size}: "
             f"got seqlen_squared_sum={seqlen_squared_sum}, expected {expected_sum_sq}"
         )
+
+
+def _make_dsv4_args():
+    """Minimal args for a DSv4-hybrid MLA model with sparse attention.
+
+    4 layers with compress_ratios [0, 4, 128, 128] (1 r0, 1 r4, 2 r128).
+    No MoE / MTP to keep the golden reference simple.
+    """
+    args = _make_gpt_args(
+        num_layers=4,
+        hidden_size=512,
+        num_attention_heads=8,
+        seq_length=256,
+        ffn_hidden_size=2048,
+        padded_vocab_size=1024,
+    )
+    args.multi_latent_attention = True
+    args.group_query_attention = False
+    args.q_lora_rank = 128
+    args.qk_head_dim = 64
+    args.qk_pos_emb_head_dim = 32
+    args.kv_lora_rank = 64
+    args.v_head_dim = 64
+    args.o_lora_rank = 64
+    args.o_groups = 2
+    args.experimental_attention_variant = "dsv4_hybrid"
+    args.csa_window_size = 64
+    args.csa_compress_ratios = [0, 4, 128, 128]
+    args.dsa_indexer_n_heads = 4
+    args.dsa_indexer_head_dim = 32
+    args.dsa_indexer_topk = 16
+    return args
+
+
+def _dsv4_golden_flops(args, total_tokens, seqlen_squared_sum):
+    """Independent golden calculator for DSv4-hybrid FLOPs.
+
+    Reimplements the formula from ``num_floating_point_operations`` so that the
+    test does not just call the same code twice. Assumes no MoE / MTP.
+    """
+    fwd_bwd = 3
+    fma = 2
+    ffn_exp = 3 if args.swiglu else 2
+
+    # ---- MLA projections (token-linear, per layer) ----
+    q_term = args.q_lora_rank * (args.hidden_size + args.num_attention_heads * args.v_head_dim + 1)
+    kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim
+    o_term = (
+        args.num_attention_heads * args.v_head_dim * args.o_lora_rank
+        + args.o_groups * args.o_lora_rank * args.hidden_size
+    )
+    mla_proj_per_layer = fwd_bwd * fma * (q_term + kv_term + o_term)
+
+    # ---- DSv4 sparse attention extra (token-linear + L^2) ----
+    ratios = args.csa_compress_ratios
+    n_r0 = sum(1 for r in ratios if r == 0)
+    n_r4 = sum(1 for r in ratios if r == 4)
+    n_r128 = sum(1 for r in ratios if r == 128)
+    nh = args.num_attention_heads
+    vhd = args.v_head_dim
+    w = args.csa_window_size
+
+    # Token-linear sparse attention
+    sparse_r0 = n_r0 * nh * w * vhd * 2
+    sparse_r128_win = n_r128 * nh * w * vhd * 2
+    if n_r4 > 0:
+        eff_topk = min(args.dsa_indexer_topk, args.seq_length // 4)
+        avg_comp_4 = eff_topk * (1 - eff_topk * 4 / (2 * args.seq_length))
+        sparse_r4 = n_r4 * nh * (w + avg_comp_4) * vhd * 2
+        idx_tok = (
+            n_r4 * args.hidden_size * (2 * args.dsa_indexer_head_dim) * 2
+            + n_r4 * args.q_lora_rank * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim
+            + n_r4 * args.hidden_size * args.dsa_indexer_n_heads
+        )
+        idx_core = n_r4 * args.dsa_indexer_n_heads * args.dsa_indexer_head_dim / 4
+    else:
+        sparse_r4, idx_tok, idx_core = 0, 0, 0
+
+    compressor = n_r4 * args.hidden_size * (2 * vhd) * 2 + n_r128 * args.hidden_size * (1 * vhd) * 2
+    dsv4_token = fwd_bwd * fma * (sparse_r0 + sparse_r4 + sparse_r128_win + compressor + idx_tok)
+    # L^2 core: r=128 compressed-KV + r=4 indexer scoring
+    r128_core = n_r128 * nh * vhd / 128
+    dsv4_core = fwd_bwd * fma * (r128_core + idx_core)
+
+    # ---- Aggregation ----
+    num_layers = args.num_layers
+    self_attn_term = mla_proj_per_layer * num_layers + dsv4_token
+    self_attn_core_term = dsv4_core  # standard core is 0 for DSv4
+
+    mlp = fwd_bwd * fma * args.hidden_size * (args.ffn_hidden_size * ffn_exp * num_layers)
+    logit = fwd_bwd * fma * args.hidden_size * args.padded_vocab_size
+
+    return total_tokens * (mlp + self_attn_term + logit) + seqlen_squared_sum * self_attn_core_term
+
+
+class TestDSv4Hybrid:
+    """DSv4 hybrid sparse-attention FLOPs against an independent golden calculator."""
+
+    def test_bshd(self):
+        """BSHD (uniform sequences) must match the golden calculator."""
+        args = _make_dsv4_args()
+        batch_size = 2
+        total_tokens = batch_size * args.seq_length
+        sum_sq = batch_size * args.seq_length**2
+
+        flops = num_floating_point_operations(args, batch_size)
+        expected = _dsv4_golden_flops(args, total_tokens, sum_sq)
+        assert flops == expected
+
+    def test_thd(self):
+        """THD (packed variable-length subsequences) must match the golden
+        calculator and be strictly less than BSHD due to the L^2 sparse-attention
+        components (r=128 compressed-KV, r=4 indexer scoring)."""
+        args = _make_dsv4_args()
+        batch_size = 2
+        packed_lengths = [64, 64, 128, 256]
+        total_tokens = sum(packed_lengths)
+        thd_sum_sq = sum(L**2 for L in packed_lengths)
+
+        flops = num_floating_point_operations(
+            args,
+            batch_size,
+            seqlen_squared_sum_in_batch=thd_sum_sq,
+            total_real_tokens_in_batch=total_tokens,
+        )
+        expected = _dsv4_golden_flops(args, total_tokens, thd_sum_sq)
+        assert flops == expected
+        # THD must be strictly less than BSHD.
+        bshd_flops = num_floating_point_operations(args, batch_size)
+        assert flops < bshd_flops