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

[niyunsheng]

Summary

num_floating_point_operations only applies sparse-attention corrections
when experimental_attention_variant == "dsv4_hybrid". The "dsa" variant (GLM-5.2) falls through to the plain multi-latent-attention branch, which counts core attention over the full causal mask and drops the indexer
entirely.

Problem

DSA scores every past position with the indexer but attends only to the
dsa_indexer_topk best keys, so the dense L^2 / 2 pair count overcounts
once sequences exceed top-k, while the indexer's own projections and its
dense O(L^2) scoring pass were never counted at all.

Measured on the full GLM-5.2 architecture (78L, hidden 6144, MoE 256E,
MTP 1, dsa_indexer_topk=2048, topk_freq=4, skip_offset=3; BSHD
uniform sequences; "before" = the plain-MLA branch the "dsa" variant
previously fell into, i.e. this function with the variant unset):

seq length	before (GFLOPs/token)	after (GFLOPs/token)	overcount
4,096	280.9	275.3	1.02x
32,768	503.6	290.0	1.74x
131,072	1,267.0	317.3	3.99x
262,144	2,284.9	352.9	6.48x

Before the fix, per-token FLOPs grow linearly with context because of the
dense L^2/2 term; after, top-k caps the core-attention term and only the
indexer's dense scoring keeps a (much smaller) per-token growth — the
physically correct shape for sparse attention. At 4K the two nearly
coincide: the mean sequence length barely exceeds top-k, so the sparsity
saving and the newly counted indexer cost roughly cancel.

Changes

megatron/training/training.py:

• Scale the core-attention coefficient by the ratio of top-k attended pairs
  to dense causal pairs (_dsa_sparse_core_scale). Short sequences
  (mean_seqlen <= topk) clamp the scale to 1.0, reducing exactly to the
  dense count — zero behavior change for non-DSA configs and for DSA runs
  within top-k.
• Add _dsa_indexer_flops: indexer projections (wq_b, wk,
  weights_proj) plus the dense O(L^2) scoring pass, charged only on
  layers that compute their own top-k index (intermediate layers reuse the
  most recent result, per is_dsa_skip_topk_layer).
• Document the accounting convention:
  • the KL-divergence indexer loss and top-k selection are excluded, like
    all other auxiliary/sorting operations in this file;
  • the figure follows the non-absorbed model-FLOPs convention — DSA
    executes the absorbed-MLA path (AbsorbedMLASelfAttention), whose core
    attention costs roughly 2x less per (query, key) pair for GLM-5.2 — so
    it is comparable across configs and attention variants, but it is NOT
    an executed-FLOPs count and must not be used directly as an MFU
    numerator.

tests/unit_tests/test_num_floating_point_operations.py:

• Add TestDSA, following the existing TestDSv4Hybrid pattern:
  • _make_dsa_args: MLA + DSA args with dsa_indexer_topk_freq=1 /
    dsa_indexer_skip_topk_offset=0, so every layer computes its own index
    and the golden reference needs no skip-layer predicate;
  • _dsa_golden_flops: independent reimplementation of the formula (the
    test does not just call the same code twice);
  • test_bshd / test_thd: exact-match both paths against the golden
    calculator and verify THD < BSHD.

Test plan

• tests/unit_tests/test_num_floating_point_operations.py on a GB200 node: 35 passed, 0 failed, 2 skipped (the two skipped classes exercise the multi-rank seqlen-stats collectives, untouched by this change).
  TestDSA.test_bshd / test_thd pass; all pre-existing classes
  (BSHD backward-compat, THD scaling, hybrid, padding removal, accumulator, DSv4-hybrid) unaffected.
• GLM-5.2 proxy integration (21L/64E, 8x GB200): with the corrected
  accounting, reported TFLOP/s/GPU is consistent across sequence lengths
  (430 @ 4K vs 382 @ 32K) instead of inflating with L^2.

Notes for reviewers

• Attended pairs = sum_i min(i, topk) over query positions; the THD path applies the same scale through the seqlen_squared_sum machinery using the batch mean sequence length.
• Indexer cost is charged only on index-computing layers; with GLM-5.2's dsa_indexer_topk_freq, intermediate layers reuse the previous index.
• No change to any other attention variant's accounting.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.