[Dev] fix: account for DSA top-k sparsity and indexer cost in FLOPs

megatron/training/training.py

@@ -384,6 +384,81 @@ def consume_seqlen_stats_in_iteration() -> Tuple[Optional[float], Optional[float
     return total_real_tokens / dedup, seqlen_squared_sum / dedup
 
 
+def _dsa_sparse_core_scale(total_real_tokens, seqlen_squared_sum, dsa_indexer_topk):
+    """Fraction of dense causal (query, key) pairs that DSA actually attends to.
+
+    A DSA layer scores every past position with the indexer but runs attention
+    only over the ``dsa_indexer_topk`` highest-scoring keys, so its core
+    attention cost is ``sum_i(min(i, topk))`` pairs per sequence instead of the
+    dense causal ``L^2 / 2``. Returns the ratio between the two, i.e. the factor
+    the dense core-attention coefficient must be scaled by.
+
+    The caller only has the batch aggregates ``sum_i(L_i)`` and
+    ``sum_i(L_i ** 2)``, not the individual sequence lengths, so the ratio is
+    evaluated at the length-weighted mean ``sum(L^2) / sum(L)``. That is exact
+    when every sequence in the batch has the same length (the usual packed-THD
+    benchmark case) and, for ragged batches, weights toward the long sequences
+    that dominate attention cost. Returns ``1.0`` when the sequences are no
+    longer than ``topk``, where top-k selects everything and attention is dense.
+    """
+    if not dsa_indexer_topk or total_real_tokens <= 0 or seqlen_squared_sum <= 0:
+        return 1.0
+    mean_seqlen = seqlen_squared_sum / total_real_tokens
+    if mean_seqlen <= dsa_indexer_topk:
+        return 1.0
+    dense_pairs = mean_seqlen * mean_seqlen / 2
+    topk_pairs = dsa_indexer_topk * mean_seqlen - dsa_indexer_topk * dsa_indexer_topk / 2
+    return topk_pairs / dense_pairs
+
+
+def _dsa_indexer_flops(
+    *, hidden_size, q_lora_rank, dsa_indexer_n_heads, dsa_indexer_head_dim, num_indexer_layers
+):
+    """Per-iteration DSA indexer FLOPs coefficients.
+
+    The indexer (``DSAIndexer``) runs only on layers that compute their own
+    top-k; layers in between reuse the most recent result (see
+    ``is_dsa_skip_topk_layer``). Per computing layer it costs:
+
+      * ``linear_wq_b``:        ``q_lora_rank -> n_heads * head_dim``
+      * ``linear_wk``:          ``hidden_size -> head_dim``
+      * ``linear_weights_proj``: ``hidden_size -> n_heads``
+      * scoring: ``einsum('sbhd,tbd->sbht')`` against every past position --
+        dense ``O(L^2)`` even though the attention that consumes it is sparse,
+        which is why this term does NOT get the top-k scaling.
+
+    The indexer KL loss (``dsa_indexer_loss_coeff``) and the top-k selection
+    itself are NOT counted: like everywhere else in this file only the model's
+    defining GEMMs enter the estimate, not auxiliary-loss or sorting work.
+
+    Returns ``(token_linear, core)`` WITHOUT the fwd+bwd (x3) or FMA (x2)
+    expansion factors -- the caller applies those. Multiply ``token_linear`` by
+    the real (unpadded) token count and ``core`` by ``sum_i(L_i ** 2)``.
+    """
+    if num_indexer_layers <= 0 or not dsa_indexer_n_heads or not dsa_indexer_head_dim:
+        return 0, 0
+    index_dim = dsa_indexer_n_heads * dsa_indexer_head_dim
+    token_linear = num_indexer_layers * (
+        q_lora_rank * index_dim  # wq_b
+        + hidden_size * dsa_indexer_head_dim  # wk
+        + hidden_size * dsa_indexer_n_heads  # weights_proj
+    )
+    # ``/2`` for the causal mask, matching the standard core-attention terms.
+    core = num_indexer_layers * index_dim / 2
+    return token_linear, core
+
+
+def _num_dsa_indexer_layers(num_layers, skip_topk_offset, topk_freq):
+    """Count layers that compute their own DSA index (the rest reuse one)."""
+    from megatron.core.transformer.experimental_attention_variant.dsa import is_dsa_skip_topk_layer
+
+    return sum(
+        1
+        for layer_number in range(1, num_layers + 1)
+        if not is_dsa_skip_topk_layer(layer_number, skip_topk_offset or 0, topk_freq or 1)
+    )
+
+
 def _dsv4_hybrid_self_attention_flops(
     *,
     hidden_size,
@@ -994,6 +1069,8 @@ def transformer_flops():
 
         dsv4_hybrid_extra_term = 0
         dsv4_hybrid_extra_core_term = 0
+        dsa_extra_term = 0
+        dsa_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):
@@ -1090,6 +1167,43 @@ def transformer_flops():
             dsv4_hybrid_extra_core_term = (
                 forward_backward_expansion_factor * fma_expansion_factor * dsv4_core_term
             )
+        elif args.experimental_attention_variant == "dsa":
+            # DSA (e.g. GLM-5.2): plain MLA projections -- so the standard
+            # token-linear term computed above still applies -- but core
+            # attention runs over the indexer's top-k keys instead of the full
+            # causal mask, and the indexer itself adds projections plus its own
+            # dense O(L^2) scoring pass. Without both corrections the estimate
+            # overcounts badly at long context (top-k caps a term that would
+            # otherwise grow with L) while silently dropping the indexer.
+            # Non-absorbed model FLOPs (same convention as the rest of this
+            # file). DSA runs the absorbed-MLA path in practice, so the sparse
+            # core term undercounts executed FLOPs by ~2x for GLM-5.2;
+            # projection terms are unaffected (absorption relocates the same
+            # W_UK/W_UV GEMMs, it does not eliminate them).
+            num_linear_attention_layers = 0
+            linear_self_attn_term = 0
+            num_standard_attention_layers = num_layers
+
+            standard_self_attn_core_term *= _dsa_sparse_core_scale(
+                total_real_tokens_in_batch, seqlen_squared_sum_in_batch, args.dsa_indexer_topk
+            )
+            indexer_token_term, indexer_core_term = _dsa_indexer_flops(
+                hidden_size=args.hidden_size,
+                q_lora_rank=(
+                    args.q_lora_rank if args.q_lora_rank is not None else args.hidden_size
+                ),
+                dsa_indexer_n_heads=args.dsa_indexer_n_heads,
+                dsa_indexer_head_dim=args.dsa_indexer_head_dim,
+                num_indexer_layers=_num_dsa_indexer_layers(
+                    num_layers, args.dsa_indexer_skip_topk_offset, args.dsa_indexer_topk_freq
+                ),
+            )
+            dsa_extra_term = (
+                forward_backward_expansion_factor * fma_expansion_factor * indexer_token_term
+            )
+            dsa_extra_core_term = (
+                forward_backward_expansion_factor * fma_expansion_factor * indexer_core_term
+            )
         else:
             num_linear_attention_layers = 0
             linear_self_attn_term = 0
@@ -1101,12 +1215,16 @@ def transformer_flops():
             linear_self_attn_term * num_linear_attention_layers
             + standard_self_attn_term * num_standard_attention_layers
             + dsv4_hybrid_extra_term
+            + dsa_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.
+        # For DSA the standard coefficient was already scaled down to the top-k
+        # pair count, and the extra term carries the indexer's dense scoring.
         self_attn_core_term = (
             standard_self_attn_core_term * num_standard_attention_layers
             + dsv4_hybrid_extra_core_term
+            + dsa_extra_core_term
         )
 
         # Token-linear FLOPs scale with the real (unpadded) token count.

tests/unit_tests/test_num_floating_point_operations.py

@@ -1038,3 +1038,130 @@ def test_hybrid_overcount_is_gone(self):
         std_flops = num_floating_point_operations(standard, batch_size)
         hyb_flops = num_floating_point_operations(hybrid, batch_size)
         assert hyb_flops == std_flops
+
+
+def _make_dsa_args():
+    """Minimal MLA + DSA args (GLM-5.2 style, small scale).
+
+    Uses ``dsa_indexer_topk_freq=1`` and ``dsa_indexer_skip_topk_offset=0``
+    so every layer computes its own top-k index -- the golden reference can
+    set ``num_indexer_layers = num_layers`` without importing the skip-layer
+    predicate from megatron.core.
+    """
+    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.kv_lora_rank = 64
+    args.qk_head_dim = 48
+    args.qk_pos_emb_head_dim = 16
+    args.v_head_dim = 64
+    args.experimental_attention_variant = "dsa"
+    args.dsa_indexer_n_heads = 4
+    args.dsa_indexer_head_dim = 32
+    args.dsa_indexer_topk = 16
+    args.dsa_indexer_topk_freq = 1
+    args.dsa_indexer_skip_topk_offset = 0
+    return args
+
+
+def _dsa_golden_flops(args, total_tokens, seqlen_squared_sum):
+    """Independent golden calculator for DSA 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,
+    ``dsa_indexer_topk_freq=1``, and ``dsa_indexer_skip_topk_offset=0``
+    (every layer computes its own index).
+    """
+    fwd_bwd = 3
+    fma = 2
+    ffn_exp = 3 if args.swiglu else 2
+    num_layers = args.num_layers
+    nh = args.num_attention_heads
+
+    # ---- MLA projections (token-linear, per layer) ----
+    q_term = args.q_lora_rank * (
+        args.hidden_size + nh * (args.qk_head_dim + args.qk_pos_emb_head_dim) + 1
+    )
+    kv_term = (
+        args.kv_lora_rank * (args.hidden_size + nh * (args.qk_head_dim + args.v_head_dim) + 1)
+        + args.hidden_size * args.qk_pos_emb_head_dim
+    )
+    o_term = nh * args.v_head_dim * args.hidden_size
+    mla_proj_per_layer = fwd_bwd * fma * (q_term + kv_term + o_term)
+
+    # ---- Core attention: MLA cost scaled down to top-k sparse pairs ----
+    raw_core = nh * (args.qk_head_dim + args.qk_pos_emb_head_dim) / 2 + nh * args.v_head_dim / 2
+    mean_seqlen = seqlen_squared_sum / total_tokens
+    topk = args.dsa_indexer_topk
+    if mean_seqlen <= topk:
+        sparse_scale = 1.0
+    else:
+        dense_pairs = mean_seqlen * mean_seqlen / 2
+        topk_pairs = topk * mean_seqlen - topk * topk / 2
+        sparse_scale = topk_pairs / dense_pairs
+    sparse_core_per_layer = fwd_bwd * fma * raw_core * sparse_scale
+
+    # ---- DSA indexer: every layer is a computing layer (topk_freq=1, skip_offset=0) ----
+    num_indexer_layers = num_layers
+    idx_dim = args.dsa_indexer_n_heads * args.dsa_indexer_head_dim
+    idx_token = num_indexer_layers * (
+        args.q_lora_rank * idx_dim  # wq_b
+        + args.hidden_size * args.dsa_indexer_head_dim  # wk
+        + args.hidden_size * args.dsa_indexer_n_heads  # weights_proj
+    )
+    idx_core = num_indexer_layers * idx_dim / 2
+    dsa_extra_token = fwd_bwd * fma * idx_token
+    dsa_extra_core = fwd_bwd * fma * idx_core
+
+    # ---- Aggregation ----
+    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
+    self_attn_term = mla_proj_per_layer * num_layers + dsa_extra_token
+    self_attn_core_term = sparse_core_per_layer * num_layers + dsa_extra_core
+
+    return total_tokens * (mlp + self_attn_term + logit) + seqlen_squared_sum * self_attn_core_term
+
+
+class TestDSA:
+    """DSA sparse-attention FLOPs against an independent golden calculator."""
+
+    def test_bshd(self):
+        """BSHD (uniform sequences) must match the golden calculator."""
+        args = _make_dsa_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 = _dsa_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 terms
+        (indexer scoring and sparse core attention)."""
+        args = _make_dsa_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 = _dsa_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