[dev] fix(moe): Fix DSA spec and rope. by yuzhongw-nvidia · Pull Request #3402 · NVIDIA/Megatron-LM

yuzhongw-nvidia · 2026-02-13T08:05:09Z

What does this PR do ?

RoPE
- Refactor RoPE utils by introducing an explicit mla_rotary_interleaved flag, keeping backward compatibility for the old argument, and adding a safe fallback to unfused RoPE when MLA-style interleaving conflicts with apply_rope_fusion.
- Disable RoPE interleaving for DSA indexer following DeepSeek's implementation code.
- Fix the PE/NoPE split-and-concat order in DSA's RoPE application path.
Experimental attention variant specs
- Disable qk layernorm fusion for DSA by forcing unfused q/kv up projections, wiring explicit q/k layernorm modules for DSA indexer input, and setting fuse_input_layernorm=False in attention metainfo.
- Fix get_experimental_attention_variant_module_spec to handle DSA dispatch.
- Add an assertion to prevent models with experimental attention variants from going through regular GPT spec functions.
- Fix MTP layer spec construction for experimental attention variants.
Functional tests
- Add a lightweight functional test case for DSA (gpt3_mcore_te_tp2_pp2_dsa) and register it in H100 recipe scopes (including slim CI scope).
Unit tests
- Add unit tests for the RoPE refactor, including MLA interleaving behavior and fusion-conflict warning/fallback coverage.
- Add comprehensive unit tests for experimental attention variant specs, covering DSA spec wiring, qk layernorm behavior, variant dispatch, attention/MoE pattern handling, and pipeline/vPP layer slicing.

PR for main: #3026, #3403

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2026-02-13T08:05:13Z

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.

yuzhongw-nvidia · 2026-02-13T08:08:01Z

/ok to test 96495c5

yuzhongw-nvidia · 2026-02-13T08:39:51Z

/ok to test 19c5f45

hxbai · 2026-02-13T09:51:03Z

        ),
    )

+    # Adjust for RMS norm.


Could you leave a comment on why layernorm_linear cannot be used here? Thanks.

Yes. I add some comments below. Does it look good?

yuzhongw-nvidia · 2026-02-25T06:51:37Z

/ok to test be5648f

yuzhongw-nvidia · 2026-02-25T08:09:23Z

/ok to test 35d6110

yuzhongw-nvidia · 2026-03-02T08:42:26Z

/ok to test f935a3b

fix uv.lock

…ndexer Update attention.py Update transformer.py Update test_rope.py rename multi_latent_attention -> mla_rotary_interleaved

yuzhongw-nvidia · 2026-03-05T09:29:19Z

/ok to test 311b567

kunlunl · 2026-03-06T02:29:07Z

/ok to test 3307414

svcnvidia-nemo-ci · 2026-03-06T07:21:46Z

🔄 Merge queue validation started!

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

yuzhongw-nvidia · 2026-03-06T08:59:33Z

/ok to test f362a99

yuzhongw-nvidia · 2026-03-06T10:04:07Z

/ok to test 484fd1a

svcnvidia-nemo-ci · 2026-03-06T12:38:07Z

🔄 Merge queue validation started!

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

yuzhongw-nvidia requested review from a team as code owners February 13, 2026 08:05

yuzhongw-nvidia mentioned this pull request Feb 13, 2026

[main] ci(moe): Add a functional test for DSA in lightweight mode #3403

Draft

6 tasks

copy-pr-bot Bot temporarily deployed to nemo-ci February 13, 2026 08:08 Inactive

ko3n1g added this to the Core 0.16 milestone Feb 13, 2026

copy-pr-bot Bot temporarily deployed to nemo-ci February 13, 2026 08:08 Inactive

copy-pr-bot Bot temporarily deployed to test February 13, 2026 08:08 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci February 13, 2026 08:39 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci February 13, 2026 08:40 Inactive

copy-pr-bot Bot temporarily deployed to test February 13, 2026 08:41 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci February 13, 2026 09:09 Inactive

hxbai reviewed Feb 13, 2026

View reviewed changes

yanring added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Feb 24, 2026

yuzhongw-nvidia force-pushed the yuzhongw/fix_dsa_spec_dev branch from 19c5f45 to be5648f Compare February 25, 2026 06:51

yuzhongw-nvidia requested a review from a team as a code owner February 25, 2026 06:51

copy-pr-bot Bot temporarily deployed to test February 25, 2026 06:52 Inactive

yuzhongw-nvidia force-pushed the yuzhongw/fix_dsa_spec_dev branch from be5648f to 35d6110 Compare February 25, 2026 08:09

copy-pr-bot Bot temporarily deployed to test February 25, 2026 08:10 Inactive

yanring self-requested a review February 26, 2026 04:40

copy-pr-bot Bot temporarily deployed to test March 2, 2026 08:43 Inactive

yuzhongw-nvidia added 8 commits March 5, 2026 01:28

fix

1069227

update dependency

8851869

fix uv.lock

add comment about why using unfused qk layernorm

54dec8e

fix rope order

ed27bd0

refactor apply_rotary_pos_emb and disable rope interleaving for DSA i…

7f8cc87

…ndexer Update attention.py Update transformer.py Update test_rope.py rename multi_latent_attention -> mla_rotary_interleaved

add a UT for apply rope refactor

e003145

add a UT for exp spec

f58b762

update exp spec UT

311b567

yuzhongw-nvidia force-pushed the yuzhongw/fix_dsa_spec_dev branch from 7324f19 to 311b567 Compare March 5, 2026 09:29

copy-pr-bot Bot temporarily deployed to test March 5, 2026 09:30 Inactive

Merge branch 'dev' into yuzhongw/fix_dsa_spec_dev

3307414

copy-pr-bot Bot temporarily deployed to test March 6, 2026 02:29 Inactive

ko3n1g approved these changes Mar 6, 2026

View reviewed changes

yanring added this pull request to the merge queue Mar 6, 2026

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

Merge branch 'dev' into yuzhongw/fix_dsa_spec_dev

f362a99

copy-pr-bot Bot had a problem deploying to test March 6, 2026 09:00 Error

hxbai enabled auto-merge March 6, 2026 09:31

fix mhc

484fd1a

copy-pr-bot Bot temporarily deployed to test March 6, 2026 10:05 Inactive

hxbai added this pull request to the merge queue Mar 6, 2026

Merged via the queue into NVIDIA:dev with commit 3d097e5 Mar 6, 2026
48 of 49 checks passed

yuzhongw-nvidia mentioned this pull request Apr 8, 2026

[QUESTION]Different RoPE Split Order for Q/K in DSA Implementation #2733

Closed

Uh oh!

Conversation

yuzhongw-nvidia commented Feb 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Feb 13, 2026

Uh oh!

yuzhongw-nvidia commented Feb 13, 2026

Uh oh!

yuzhongw-nvidia commented Feb 13, 2026

Uh oh!

hxbai Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia Feb 26, 2026

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia commented Feb 25, 2026

Uh oh!

yuzhongw-nvidia commented Feb 25, 2026

Uh oh!

yuzhongw-nvidia commented Mar 2, 2026

Uh oh!

yuzhongw-nvidia commented Mar 5, 2026

Uh oh!

kunlunl commented Mar 6, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 6, 2026

Uh oh!

Uh oh!

yuzhongw-nvidia commented Mar 6, 2026

Uh oh!

yuzhongw-nvidia commented Mar 6, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

yuzhongw-nvidia commented Feb 13, 2026 •

edited

Loading

(Step 1): Add PR label `Expert Review`