fix(training): align DSv4 packed FLOPs with MCore

[cuichenx]

What does this PR do?

Align DeepSeek-V4 packed-sequence FLOPs accounting with Megatron-LM's token-linear and quadratic split.

Changelog

• Scale ratio-128 compressed-KV attention and ratio-4 indexer scoring with the runtime seqlen_squared_sum.
• Keep ratio-4 nominal top-k selection tied to the configured sequence length, matching Megatron-LM.
• Add an independent packed BSHD/THD regression that checks the Megatron-LM golden expression and verifies lower THD FLOPs at equal token count.

Root cause and impact

The Bridge formula predated Megatron-LM's packed-sequence correction in NVIDIA/Megatron-LM#5358. It collapsed all DSv4 sparse-attention work into a token-linear expression using one effective sequence length. That is equivalent for the supported fixed-length cases covered by the existing test, but diverges for packed variable-length sequences.

This change preserves supported fixed-length results while making packed DSv4 accounting match Megatron-LM:

F = total_tokens * token_linear_term + seqlen_squared_sum * quadratic_term

No router, mHC, MTP, or backend-specific accounting is changed because those are shared calculator omissions rather than Bridge/Megatron-LM parity mismatches.

GitHub Actions CI

No GPU or cluster jobs were run for this CPU-only FLOPs change.

Validation

• uv run python -m pytest tests/unit_tests/training/utils/test_flop_utils.py::TestDeepSeekV4HybridFlops -q — 10 passed in a dependency-isolated CPU environment.
• uv run pre-commit run --all-files — passed.
• Independent Megatron-LM golden values matched for both BSHD (32,548,061,184) and THD (32,538,230,784) fixtures.

Before your PR is "Ready for review"

• Read and followed the contributor guidelines.
• Added the necessary unit test.
• No documentation update is required for this parity fix.
• This PR does not affect optional-install components.

Additional Information

• Megatron-LM reference fix: [Dev] restore DSv4 tflops calc in training and fix the packed seq case NVIDIA/Megatron-LM#5358

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cuichenx]

/claude review

[claude]

LGTM - light review, one observation.

Verified the DSv4 packed-FLOPs split analytically:

• The compressed-KV (sparse_attn_r128_core) and indexer-scoring (indexer_scoring_core) terms are correctly moved out of the token-linear self_attn_term into dsv4_hybrid_core_attn_term, scaled by effective_seqlen_squared_sum. The old avg_comp_128 = (... // 128)/2 with the trailing *2 cancels exactly against the new /128 core term, so the split is algebraically faithful.
• dsv4_hybrid_core_attn_term is initialized to 0, so every non-DSv4 path adds effective_seqlen_squared_sum * 0 - no regression for other models.
• The unchanged test_dsv4_hybrid_exact_flops (seq_len=512) still passes because 512 is divisible by 128 and 4, making the BSHD fallback (effective_seqlen_squared_sum = seqlen_sum * effective_seq_length) reproduce the legacy value.

Observation (not blocking): legacy BSHD equivalence holds only when seq_length is a multiple of 128 (and 4). The old code used integer floor // 128; the new quadratic term uses true division, so for a seq_length not divisible by 128 the fixed-length result now differs slightly from the pre-PR value. This is the intended MCore-parity direction, just worth noting for anyone comparing against older numbers.

Note: I could not execute the unit tests in the review sandbox (pytest was blocked); the assessment above is by inspection. The PR reports 10 passed locally.

Suggested test cases:

• No perf tests impacted. The change is confined to src/megatron/bridge/training/utils/flop_utils.py and its unit test; no files under scripts/performance/configs/ are touched. Relevant unit coverage: tests/unit_tests/training/utils/test_flop_utils.py::TestDeepSeekV4HybridFlops::test_dsv4_hybrid_exact_flops and ::test_dsv4_hybrid_packed_flops_match_mcore_split.