Account for DSA in model FLOPs

megatron/training/training.py

@@ -777,14 +777,188 @@ def transformer_flops():
             linear_self_attn_term = 0
             num_standard_attention_layers = num_layers
 
-        # Token-linear self-attention work (projections + linear-attention variants).
-        # Linear attention has no L^2 term, so it stays entirely token-linear.
-        self_attn_term = (
-            linear_self_attn_term * num_linear_attention_layers
-            + standard_self_attn_term * num_standard_attention_layers
-        )
-        # Core attention (L^2) FLOPs per standard-attention layer.
-        self_attn_core_term = standard_self_attn_core_term * num_standard_attention_layers
+        if args.experimental_attention_variant == "dsa":
+            # DSA replaces dense MLA core attention with top-k attention while retaining a
+            # dense lightning indexer. The attention/indexer geometry follows equations 1-2
+            # of the official report: https://arxiv.org/abs/2512.02556. MCore implements the
+            # MLA path with matrix absorption: QK operates over kv_lora_rank + RoPE channels
+            # and AV over kv_lora_rank channels (experimental_attention_variant/absorbed_mla.py).
+            if not args.multi_latent_attention:
+                raise ValueError("dsa FLOPs calculation requires multi_latent_attention")
+
+            qk_head_dim = getattr(args, "qk_head_dim", 64)
+            qk_pos_emb_head_dim = getattr(args, "qk_pos_emb_head_dim", 0)
+            v_head_dim = getattr(args, "v_head_dim", 64)
+            kv_lora_rank = getattr(args, "kv_lora_rank", 0)
+            if kv_lora_rank is None or kv_lora_rank <= 0:
+                raise ValueError("kv_lora_rank must be positive for dsa FLOPs calculation")
+
+            idx_n_heads = getattr(args, "dsa_indexer_n_heads", None)
+            idx_head_dim = getattr(args, "dsa_indexer_head_dim", None)
+            idx_topk = getattr(args, "dsa_indexer_topk", None)
+            if idx_n_heads is None or idx_n_heads <= 0:
+                raise ValueError("dsa_indexer_n_heads must be positive for dsa FLOPs calculation")
+            if idx_head_dim is None or idx_head_dim <= 0:
+                raise ValueError("dsa_indexer_head_dim must be positive for dsa FLOPs calculation")
+            if idx_topk is None or idx_topk <= 0:
+                raise ValueError("dsa_indexer_topk must be positive for dsa FLOPs calculation")
+
+            idx_topk_freq = getattr(args, "dsa_indexer_topk_freq", 1)
+            idx_skip_topk_offset = getattr(args, "dsa_indexer_skip_topk_offset", 0)
+            if not isinstance(idx_topk_freq, int) or idx_topk_freq <= 0:
+                raise ValueError("dsa_indexer_topk_freq must be a positive integer")
+            if not isinstance(idx_skip_topk_offset, int) or idx_skip_topk_offset < 0:
+                raise ValueError("dsa_indexer_skip_topk_offset must be a non-negative integer")
+
+            def count_indexer_layers(layer_count: int) -> int:
+                """Count layers that compute rather than reuse DSA top-k indices."""
+                sharing_offset = max(idx_skip_topk_offset, 1)
+                return sum(
+                    max(layer_number - sharing_offset, 0) % idx_topk_freq == 0
+                    for layer_number in range(1, layer_count + 1)
+                )
+
+            # MCore gives MTP layers their own 1-based numbering, so their sharing cadence
+            # restarts independently of the decoder. GLM-5.2 has MTP1, which computes one
+            # fresh index. See transformer/multi_token_prediction.py in Megatron-Core.
+            num_indexer_layers = count_indexer_layers(args.num_layers) + count_indexer_layers(
+                mtp_num_layers
+            )
+
+            # Average valid causal pairs per token. The top-k expression is exact for a
+            # fixed-length batch. Packed metadata supplies only sum(s) and sum(s^2), so for
+            # mixed sequences crossing top-k we use the token-weighted effective length; it
+            # remains exact when every subsequence is <= top-k.
+            if total_real_tokens_in_batch > 0:
+                core_attn_seq_factor = (
+                    seqlen_squared_sum_in_batch / total_real_tokens_in_batch
+                )
+            else:
+                core_attn_seq_factor = args.seq_length
+            dense_causal_context = (core_attn_seq_factor + 1) / 2
+            if core_attn_seq_factor <= idx_topk:
+                sparse_causal_context = dense_causal_context
+            else:
+                sparse_causal_context = idx_topk - idx_topk * (idx_topk - 1) / (
+                    2 * core_attn_seq_factor
+                )
+
+            if args.q_lora_rank is None:
+                q_term = (
+                    args.hidden_size
+                    * args.num_attention_heads
+                    * (qk_head_dim + qk_pos_emb_head_dim)
+                )
+                indexer_q_input_size = args.hidden_size
+            else:
+                q_term = args.q_lora_rank * (
+                    args.hidden_size
+                    + args.num_attention_heads * (qk_head_dim + qk_pos_emb_head_dim)
+                    + 1  # q norm
+                )
+                indexer_q_input_size = args.q_lora_rank
+
+            kv_term = (
+                kv_lora_rank
+                * (
+                    args.hidden_size
+                    + args.num_attention_heads * (qk_head_dim + v_head_dim)
+                    + 1  # kv norm
+                )
+                + args.hidden_size * qk_pos_emb_head_dim
+            )
+            output_term = args.num_attention_heads * v_head_dim * args.hidden_size
+            mla_projection_term = (
+                forward_backward_expansion_factor
+                * fma_expansion_factor
+                * num_layers
+                * (q_term + kv_term + output_term)
+            )
+
+            absorbed_qk_dim = kv_lora_rank + qk_pos_emb_head_dim
+            absorbed_v_dim = kv_lora_rank
+            sparse_mla_core_term = (
+                forward_backward_expansion_factor
+                * fma_expansion_factor
+                * num_layers
+                * sparse_causal_context
+                * args.num_attention_heads
+                * (absorbed_qk_dim + absorbed_v_dim)
+            )
+
+            indexer_projection_size = (
+                indexer_q_input_size * idx_n_heads * idx_head_dim
+                + args.hidden_size * idx_head_dim
+                + args.hidden_size * idx_n_heads
+            )
+            indexer_loss_coeff = getattr(args, "dsa_indexer_loss_coeff", 0.0) or 0.0
+            trains_indexer = indexer_loss_coeff > 0
+
+            # MCore detaches x and q_resid before the indexer. With the auxiliary loss
+            # enabled, each projection therefore executes forward+wgrad (2x forward), not
+            # forward+dgrad+wgrad (3x). Without the loss, top-k is computed under no_grad.
+            indexer_projection_multiplier = 4 if trains_indexer else 2
+            indexer_projection_term = (
+                indexer_projection_multiplier * num_indexer_layers * indexer_projection_size
+            )
+
+            # Equation 1 computes H_i dense q_i.k_i dot products and a weighted head
+            # reduction. Top-k selection needs every causal score, so the forward is
+            # always dense. The backward only covers score entries the KL loss touches:
+            # every causal pair for the dense loss, but only the selected top-k pairs for
+            # the sparse loss (no gradient flows through the discrete top-k selection;
+            # MCore's indexer_backward_wrapper consumes the selected payload only).
+            # ReLU, normalization, and top-k comparisons are not floating-point matmuls and
+            # are intentionally outside this model-FLOPs numerator.
+            use_sparse_indexer_loss = getattr(args, "dsa_indexer_use_sparse_loss", False)
+            index_score_unit = idx_n_heads * (idx_head_dim + 1)
+            index_score_term = (
+                2 * num_indexer_layers * dense_causal_context * index_score_unit
+            )
+            if trains_indexer:
+                score_grad_context = (
+                    sparse_causal_context if use_sparse_indexer_loss else dense_causal_context
+                )
+                index_score_term += (
+                    4 * num_indexer_layers * score_grad_context * index_score_unit
+                )
+
+            # GLM recipes enable MCore's sparse indexer KL loss. Its attention target uses
+            # detached main-model Q/K, hence forward-only QK work. Dense-loss configurations
+            # use the full causal context instead of the selected context.
+            indexer_teacher_term = 0
+            if trains_indexer:
+                teacher_context = (
+                    sparse_causal_context if use_sparse_indexer_loss else dense_causal_context
+                )
+                indexer_teacher_term = (
+                    2
+                    * num_indexer_layers
+                    * teacher_context
+                    * args.num_attention_heads
+                    * absorbed_qk_dim
+                )
+
+            # All DSA pairwise work is expressed as average causal context per real token.
+            # This keeps the exact dense packed-sequence contribution while sharing the
+            # same effective-length approximation as Bridge for sparse packed sequences.
+            self_attn_term = (
+                mla_projection_term
+                + sparse_mla_core_term
+                + indexer_projection_term
+                + index_score_term
+                + indexer_teacher_term
+            )
+            self_attn_core_term = 0
+        else:
+            # Token-linear self-attention work (projections + linear-attention variants).
+            # Linear attention has no L^2 term, so it stays entirely token-linear.
+            self_attn_term = (
+                linear_self_attn_term * num_linear_attention_layers
+                + standard_self_attn_term * num_standard_attention_layers
+            )
+            # 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

@@ -99,6 +99,37 @@ def _make_hybrid_args(*, num_layers=4, hidden_size=512, num_attention_heads=8, s
     return args
 
 
+def _make_dsa_args(**overrides):
+    """Minimal args for absorbed MLA with Dynamic Sparse Attention."""
+    values = {
+        "num_layers": 1,
+        "hidden_size": 8,
+        "num_attention_heads": 2,
+        "seq_length": 4,
+        "swiglu": False,
+        "ffn_hidden_size": 0,
+        "padded_vocab_size": 128,
+    }
+    values.update(overrides)
+    args = _make_gpt_args(**values)
+    args.multi_latent_attention = True
+    args.group_query_attention = False
+    args.experimental_attention_variant = "dsa"
+    args.q_lora_rank = 4
+    args.kv_lora_rank = 3
+    args.qk_head_dim = 2
+    args.qk_pos_emb_head_dim = 1
+    args.v_head_dim = 2
+    args.dsa_indexer_n_heads = 2
+    args.dsa_indexer_head_dim = 2
+    args.dsa_indexer_topk = 2
+    args.dsa_indexer_topk_freq = 1
+    args.dsa_indexer_skip_topk_offset = 0
+    args.dsa_indexer_loss_coeff = 0.001
+    args.dsa_indexer_use_sparse_loss = True
+    return args
+
+
 class TestBSHDBackwardCompat:
     """For unpacked BSHD, the new optional arg must not change the result."""
 
@@ -150,6 +181,83 @@ def test_mla_default_matches_explicit_bshd(self):
         assert default_flops == explicit_flops
 
 
+class TestDynamicSparseAttentionFlops:
+    """Closed-form training FLOPs for absorbed MLA with Dynamic Sparse Attention."""
+
+    def test_dsa_exact_toy_formula(self):
+        """A one-layer toy covers absorbed sparse MLA and every lightning-indexer matmul."""
+        # At S=4 and top-k=2, the causal selected counts are [1, 2, 2, 2],
+        # while the dense indexer sees [1, 2, 3, 4]. The independently reduced
+        # per-token terms are:
+        #   MLA projections: 906; sparse QK+AV: 147
+        #   indexer projections (forward+wgrad): 192
+        #   index scores (dense forward 30 + sparse-loss backward 42): 72
+        #   sparse detached teacher QK (forward only): 28
+        #   vocabulary projection: 6144
+        expected = 4 * (906 + 147 + 192 + 72 + 28 + 6144)
+
+        assert num_floating_point_operations(_make_dsa_args(), batch_size=1) == expected
+
+    def test_dsa_sequence_length_and_topk_scaling(self):
+        """Sparse MLA saturates at top-k while the lightning indexer remains quadratic."""
+        short = num_floating_point_operations(_make_dsa_args(), batch_size=1)
+        long = num_floating_point_operations(_make_dsa_args(seq_length=8), batch_size=1)
+        wider_topk_args = _make_dsa_args()
+        wider_topk_args.dsa_indexer_topk = 4
+        wider_topk = num_floating_point_operations(wider_topk_args, batch_size=1)
+
+        # S=8, k=2: average sparse context is 15/8 and dense causal context is 9/2.
+        assert long == 60_228
+        assert long > 2 * short
+        # S=4, k=4 raises average sparse context from 7/4 to 5/2. Sparse QK/AV,
+        # the sparse-loss score backward, and the sparse teacher target change;
+        # top-k comparisons are not FLOPs.
+        assert wider_topk - short == 372
+
+    def test_dsa_index_sharing_cadence_and_offset(self):
+        """Only full IndexShare layers pay indexer projection, score, and teacher work."""
+        no_sharing_args = _make_dsa_args(num_layers=6)
+        no_sharing_args.dsa_indexer_topk_freq = 1
+        no_sharing = num_floating_point_operations(no_sharing_args, batch_size=1)
+
+        offset_three_args = _make_dsa_args(num_layers=6)
+        offset_three_args.dsa_indexer_topk_freq = 4
+        offset_three_args.dsa_indexer_skip_topk_offset = 3
+        offset_three = num_floating_point_operations(offset_three_args, batch_size=1)
+
+        offset_one_args = _make_dsa_args(num_layers=6)
+        offset_one_args.dsa_indexer_topk_freq = 4
+        offset_one_args.dsa_indexer_skip_topk_offset = 1
+        offset_one = num_floating_point_operations(offset_one_args, batch_size=1)
+
+        # Full layers are [1..6], [1,2,3], and [1,5], respectively. Each full
+        # layer contributes (192 + 72 + 28) FLOPs per token of indexer work.
+        assert no_sharing - offset_three == 3 * 4 * 292
+        assert offset_three - offset_one == 4 * 292
+
+    def test_dsa_detached_indexer_projection_backward_multiplier(self):
+        """Indexer loss adds wgrad, not dgrad, for projections fed by detached inputs."""
+        with_loss = num_floating_point_operations(_make_dsa_args(), batch_size=1)
+        without_loss_args = _make_dsa_args()
+        without_loss_args.dsa_indexer_loss_coeff = 0.0
+        without_loss = num_floating_point_operations(without_loss_args, batch_size=1)
+
+        # Enabling sparse indexer loss adds one projection wgrad (96/token),
+        # score gradients over the selected top-k context only (42/token),
+        # and teacher QK (28/token).
+        assert with_loss - without_loss == 4 * (96 + 42 + 28)
+
+    def test_dsa_mtp_layer_has_independent_sparse_attention_and_indexer(self):
+        """MTP1 adds one full DSA layer because MCore restarts MTP layer numbering at one."""
+        decoder_only = num_floating_point_operations(_make_dsa_args(), batch_size=1)
+        with_mtp_args = _make_dsa_args()
+        with_mtp_args.mtp_num_layers = 1
+        with_mtp = num_floating_point_operations(with_mtp_args, batch_size=1)
+
+        # Added per-token work: DSA layer 1345 + MTP norms/eh-proj 912 + logits 6144.
+        assert with_mtp - decoder_only == 4 * (1345 + 912 + 6144)
+
+
 class TestTHDScaling:
     """Only the L^2 attention term should depend on ``seqlen_squared_sum_in_batch``."""