fix: account for DSA top-k sparsity and indexer cost in FLOPs by niyunsheng · Pull Request #7043 · NVIDIA/Megatron-LM

niyunsheng · 2026-09-02T09:54:42Z

Port the DSA FLOPs accounting from #6753 and #7157 to main for standard GPT and hybrid models. Previously, hybrid D layers contributed no attention/indexer FLOPs, while the standard path used ordinary MLA attention costs.

Changes

Share explicit-parameter counters in megatron/training/logging/dsa_flops.py. dsa_layer_flops() returns one attention layer's forward FLOPs; callers multiply by the DSA layer count and apply the common x3 training expansion.
Count indexers separately using their loss-dependent backward factors: forward only with loss disabled; projection x2 and scoring x3 for dense KL, or scoring 1 + 2 * sparse_scale for sparse KL.
dsa_indexer_flops(num_dsa_layers=...) accepts all DSA attention layers, including MTP and layers reusing an index. It derives actual indexer executions from sharing settings and layer positions. Decoder and MTP numbering are handled separately; repeated MTP counts each execution. Hybrid models use pattern positions, including numbering across non-DSA layers.
Support kpool > 1: gate projection, complete-pool key reduction, compressed causal scoring, and attention over selected pools plus the optional incomplete tail. KPool requires disabled indexer loss, matching runtime. Attention and indexer share the pool-pair sum; indexer arithmetic stays in one helper.

Changes are limited to the DSA logging module, training integration and FLOPs tests. Existing non-DSA counters and tests are preserved.

Validation

199 CPU FLOPs tests passed; 20 CUDA accumulator tests deselected. Includes independent goldens, KPool sizes 2/4/8, tail modes, sharing, MTP, fractional lengths and ten comparisons with runtime KPool token selection.
41,640 exact numerical comparisons against a4385191a validate the final helper simplification, plus 32 matching validation/zero-layer outcomes. Formatting, lint, copyright and environment-aware type checks pass.
OCI-HSG, 4 x GB200, BF16, TP2/DP2, unfused DSA: four KPool configurations at a4385191a each completed 5/5 optimizer steps: standard pool=4 with tail on/off; hybrid with two MTP depths and sharing, pool=4/tail on and pool=8/tail off. All losses/norms were finite, with zero skipped/NaN iterations. Eight broader standard/hybrid/MTP smokes passed at 6b46fe8ad.

The final helper simplification has exact CPU parity with the GPU-tested implementation; GPU training was not rerun after it. The squash commit preserves the exact file tree of 914c37679. GPU coverage does not include fused cuDNN DSA, packed THD training or eight-GPU accumulator topologies.

Accounting limits

Packed batches expose only sum(L) and sum(L^2): KPool uses the length-weighted mean with interpolation, so individual document/pool boundaries remain approximate. Pool softmax, positional bias, top-k selection, teacher recomputation and KL operations are excluded. The pre-existing ordinary DSA short-THD approximation and hybrid MTP norm/projection gap are unchanged.

copy-pr-bot · 2026-09-02T09:54:46Z

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.

Port of NVIDIA#6753 (dev) onto main, extended for the hybrid path. Unlike dev, main builds DSA models only through the hybrid path (--hybrid-layer-pattern with 'D' layers; the standard GPT path rejects experimental attention variants), so instead of dev's fail-loud assert this adds hybrid 'D'-layer accounting: absorbed-MLA projections, top-k sparse core attention, and the indexer counted on block-local layer numbers (main pattern sequential across pipes, MTP depths restart at 1). (cherry picked from commit 85902ef) Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yunsheng Ni <yunshengn@nvidia.com>

niyunsheng · 2026-09-02T12:25:06Z

/ok to test 7f5d522

With `--dsa-indexer-use-sparse-loss` the indexer KL loss is taken over the selected top-k keys only, so the score gradient is zero outside them and dq/dk of the indexer scoring GEMM span just the top-k pairs (this is exactly what the fused cuDNN backend executes). The previous commit always charged the dense fwd + dq + dk = 3x, overcounting the indexer's O(L^2) term by up to ~3x at long context when the sparse loss is enabled. * `_dsa_indexer_flops` now takes `dsa_indexer_use_sparse_loss` and the top-k pair ratio (`sparse_core_scale`), and charges the scoring term `1 + 2 * sparse_core_scale` for the sparse KL variant, 3x for the dense one, 1x when the indexer loss is off. Parameters are named after the corresponding args (`dsa_indexer_loss_coeff`, `dsa_indexer_use_sparse_loss`). * Both the standard-path and the hybrid-path ('D' layers) call sites pass the flag and the shared `_dsa_sparse_core_scale`. * Golden calculator in the unit test follows the same rule; new tests cover sparse vs dense parity at L <= topk, strict shrinkage past topk on both paths, the flag being a no-op without the indexer loss, and the raw expansion factors of `_dsa_indexer_flops`. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yunsheng Ni <yunshengn@nvidia.com>

copy-pr-bot · 2026-09-09T05:29:23Z

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.

niyunsheng · 2026-09-09T05:34:06Z

/ok to test 56ebb2b

deepakn94 · 2026-09-10T23:35:12Z

    return total_real_tokens / dedup, seqlen_squared_sum / dedup


+def _dsa_sparse_core_scale(total_real_tokens, seqlen_squared_sum, dsa_indexer_topk):


Can we move these new FLOPs functions and the existing ones into a separate new file in megatron/training? These functions are getting too long and megatron/training/training.py needs to be cleaned.

@deepakn94 Agreed that the FLOPs calculations should live in a separate module. Since this PR fixes DSA FLOPs accounting, would you prefer a separate refactoring PR to move all FLOPs functions into megatron/training/flops.py? That would keep the correctness changes easier to review. I’m happy to do the extraction as a follow-up, or land it first and rebase this PR if you prefer.

How about these new functions go into megatron/training/flops/dsa.py, then a follow up can move the rest into megatron/training/flops/whatever.py?

Can we use the newly created megatron/training/logging/? This directory currently only contains packed_sequence_stats.py. Alternatively, we can rename logging to stats or flops.

@jaredcasper @Phlip79 I’m happy with either location. Would you prefer training/flops/dsa.py or training/logging/dsa_flops.py? I’ll keep this PR focused on extracting the DSA helpers and leave the remaining FLOPs extraction for a follow-up.

Move DSA FLOPs accounting into megatron/training/logging/dsa_flops.py. Share a single-layer forward attention counter between standard and hybrid models, aggregate attention with the common x3 training expansion, and account for indexers separately using their executed layer count and loss-dependent backward cost. Preserve standard/hybrid indexer numbering, MTP behavior, non-DSA accounting, and the requested hybrid docstring. Validation: 45 CPU FLOPs tests passed; 23,040 exact parent comparisons and 576 large-model comparisons passed (DSA reassociation differences at most two ULPs). Four real OCI-HSG 4xGB200 TP2/DP2 BF16 training configurations each completed 5/5 iterations with finite loss/gradient norms and zero skipped/NaN iterations. GPU coverage uses unfused attention. Formatting, lint, copyright, and type checks passed. Signed-off-by: Yunsheng Ni <yunshengn@gmail.com>

Use explicit attention/indexer parameters and count indexer executions with decoder/MTP numbering and index sharing. Add KPool gate, compression, scoring and tail-attention costs. Consolidate the indexer formulas and share causal pool-pair counting while preserving existing non-DSA behavior. Squashes the six follow-up commits after fdedf45; the complete file tree is identical to 914c376. Validation: 199 CPU FLOPs tests, 41640 exact numerical comparisons and 32 validation outcomes passed; lint/type/scope checks passed. Four KPool GPU smokes passed on OCI-HSG 4xGB200 at a438519 before the equivalent helper simplification; eight broader smokes passed at 6b46fe8. Signed-off-by: Yunsheng Ni <yunshengn@gmail.com>

niyunsheng · 2026-10-26T11:53:10Z

/ok to test def7ec1

nemo-automation-bot · 2026-10-26T13:55:47Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40334912120

nemo-automation-bot · 2026-10-26T15:21:19Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40348478286

- Main's hybrid packed-schedule test (NVIDIA#7047) and SP-embedding test (NVIDIA#8274) call HybridModel.forward on stubs. Bind the _preprocess/_postprocess hooks it now delegates to, and assert on the unwrapped decoder output, since the stub decoder unwraps its WrappedTensor like HybridStack does. - Main's TransformerLayerSchedulePlan.run reads recompute_segment (NVIDIA#6429), so the schedule-run test sets it on its __new__-built plan. - Main's DSA indexer FLOPs count (NVIDIA#7043) numbered every pattern character as a layer, so [ and ] shifted the top-k skip positions. Count the flattened symbols, which keep the ungrouped numbering. - Drop an unused test import and a comment about a chunkwise-CP padding-mask check that NVIDIA#7047 removed. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

copy-pr-bot Bot temporarily deployed to test September 2, 2026 10:01 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 10:03 Inactive

niyunsheng marked this pull request as ready for review September 2, 2026 10:25

niyunsheng requested a review from a team as a code owner September 2, 2026 10:25

niyunsheng requested a review from a team September 2, 2026 10:25

github-actions Bot added the complexity: low label Sep 2, 2026

niyunsheng marked this pull request as draft September 2, 2026 10:43

niyunsheng force-pushed the yunshengn/fix-dsa-flops-accounting-main branch from 33b7228 to 7f5d522 Compare September 2, 2026 10:55

niyunsheng marked this pull request as ready for review September 2, 2026 12:24

copy-pr-bot Bot temporarily deployed to test September 2, 2026 12:26 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 12:28 Inactive

github-actions Bot added complexity: medium and removed complexity: low labels Sep 2, 2026

zhongbozhu approved these changes Sep 3, 2026

View reviewed changes

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

copy-pr-bot Bot temporarily deployed to test September 9, 2026 05:35 Inactive

copy-pr-bot Bot temporarily deployed to public September 9, 2026 05:38 Inactive

deepakn94 reviewed Sep 10, 2026

View reviewed changes

zhongbozhu enabled auto-merge October 23, 2026 17:55

niyunsheng force-pushed the yunshengn/fix-dsa-flops-accounting-main branch from 914c376 to fdedf45 Compare October 26, 2026 11:27

copy-pr-bot Bot temporarily deployed to test October 26, 2026 11:36 Inactive

copy-pr-bot Bot temporarily deployed to public October 26, 2026 11:39 Inactive

niyunsheng force-pushed the yunshengn/fix-dsa-flops-accounting-main branch from 93cdbe5 to 6d66c2a Compare October 26, 2026 11:43

niyunsheng added 2 commits October 26, 2026 19:43

Merge branch 'main' into yunshengn/fix-dsa-flops-accounting-main

def7ec1

copy-pr-bot Bot temporarily deployed to test October 26, 2026 11:55 Inactive

copy-pr-bot Bot temporarily deployed to public October 26, 2026 11:57 Inactive

zhongbozhu added this pull request to the merge queue Oct 26, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 26, 2026

niyunsheng added this pull request to the merge queue Oct 26, 2026

Merged via the queue into NVIDIA:main with commit 7052740 Oct 26, 2026
101 checks passed

niyunsheng deleted the yunshengn/fix-dsa-flops-accounting-main branch October 26, 2026 16:43

FDecaYed mentioned this pull request Oct 26, 2026

Count DSv4 attention FLOPs in HybridModel training metrics #7972

Open

1 task

kevalmorabia97 mentioned this pull request Oct 27, 2026

Add mHC precision controls, mean output contract, and full activation recompute with mHC #7521

Merged

4 of 7 tasks

BestJuly mentioned this pull request Oct 27, 2026

Qwen3.8-Flash-Next: gated-residual mHC, Qwen Sparse Attention, hashed n-gram memory, and mHC pipeline support #7553

Draft

FDecaYed mentioned this pull request Oct 28, 2026

chore: AUT-2655 nightly sync main into dev (2026-10-27) #8643

Merged

Connor-XY mentioned this pull request Oct 28, 2026

[feat] HybridStack grouped syntax and EP-overlap (2/4 of #4798) #4942

Open

		return total_real_tokens / dedup, seqlen_squared_sum / dedup


		def _dsa_sparse_core_scale(total_real_tokens, seqlen_squared_sum, dsa_indexer_topk):

Conversation

niyunsheng commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Validation

Accounting limits

Uh oh!

copy-pr-bot Bot commented Sep 2, 2026

Uh oh!

niyunsheng commented Sep 2, 2026

Uh oh!

copy-pr-bot Bot commented Sep 9, 2026

Uh oh!

niyunsheng commented Sep 9, 2026

Uh oh!

deepakn94 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

niyunsheng Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

jaredcasper Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

Phlip79 Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

niyunsheng Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

niyunsheng commented Oct 26, 2026

Uh oh!

nemo-automation-bot Bot commented Oct 26, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Oct 26, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

niyunsheng commented Sep 2, 2026 •

edited

Loading

Phlip79 Sep 23, 2026 •

edited

Loading