[main] fix(moe): Fix several bugs for DSA rope and spec.

[yuzhongw-nvidia]

What does this PR do ?

⚠️ 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.

Changes are listed as below.

• 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 dev: #3402

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]

Loading

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

1. Attach the Expert Review label when your PR is ready for review.
2. 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

1. Add Final Review label
2. 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]

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]

/ok to test f1316d1

[yuzhongw-nvidia]

/ok to test 510a86f

[Phlip79]

We are changing our review process and marking all open, unlabeled PRs as draft. This change will go in effect starting once #3659 is merged.

Moving forward, all PRs will be required to start as draft PRs. If you wish to get your PR merged, mark your PR as “Ready for review”. Read more about the new process at submit.md.

[Phlip79]

/claude review

[yuzhongw-nvidia]

/ok to test 61a46bb

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[Phlip79]

Merge is failing due to a TypeError: get_gpt_decoder_layer_specs() got an unexpected keyword argument 'vp_stage' error.

[yuzhongw-nvidia]

/ok to test 8477131

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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