Account for DSA in model FLOPs

[yaoyu-33]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Accounts for Dynamic Sparse Attention (DSA) in the training model-FLOPs numerator instead of falling through to dense MLA accounting.

This mirrors NVIDIA-NeMo/Megatron-Bridge#5324 into Megatron-LM's num_floating_point_operations implementation and adapts it to the current token-linear / packed-sequence FLOPs split.

The DSA branch includes:

• absorbed sparse MLA QK/AV work using the configured top-k;
• lightning-indexer projections, dense score computation, and weighted head reduction;
• dense or sparse indexer-loss backward work;
• detached teacher QK work;
• IndexShare frequency/offset accounting, including the independently numbered MTP stack.

For packed sequences, dense causal pairs remain exact from sum(L) and sum(L^2). Sparse pairs use the same token-weighted effective-length approximation documented in the Bridge change because the existing metadata cannot recover each individual sequence length.

Issue tracking

Linked issue: N/A — small bug fix mirrored from NVIDIA-NeMo/Megatron-Bridge#5324.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests — not needed for calculator-only arithmetic
• I have added proper typing to my code
• I have added relevant documentation — behavior is documented inline
• I have run the autoformatter — the pinned CUDA dependency set cannot resolve on macOS arm64

Validation

• Five standalone closed-form DSA checks passed: exact toy total, sequence/top-k scaling, IndexShare cadence/offset, detached-indexer backward multiplier, and MTP1.
• uvx ruff check megatron/training/training.py tests/unit_tests/test_num_floating_point_operations.py
• uvx isort --check-only megatron/training/training.py tests/unit_tests/test_num_floating_point_operations.py
• python3 -m py_compile megatron/training/training.py tests/unit_tests/test_num_floating_point_operations.py
• git diff --check

The repository's uv environment and distributed pytest suite could not run on this Apple Silicon host because nvidia-cudnn-frontend==1.26.0 has no macOS wheel. CI should run the focused unit tests in the supported CUDA environment.

[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.