[main] [DeepSeek-v4] Hybrid Attention with CSA and HCA

[hxbai]

What does this PR do ?

dev PR: #4458

Add DeepSeek-v4 Hybrid Attention with CSA and HCA.

We reuse some arguments from DSA to avoid duplication.

DeepSeek-V4 training has several phases, and here are the setting examples (use DSv4-Flash size):

• Dense mode:
  • --experimental-attention-variant dsv4_hybrid
  • --csa-window-size 128
  • --csa-compress-ratios ([0,0,4]+[128,4]*20+[0])
  • --csa-compress-rotary-base: 40000
  • --dsa-indexer-n-heads 64
  • --dsa-indexer-head-dim 128
  • --dsa-indexer-topk 512
  • --csa-dense-mode
• Sparse warmup mode:
  • --experimental-attention-variant dsv4_hybrid
  • --csa-window-size 128
  • --csa-compress-ratios ([0,0,4]+[128,4]*20+[0])
  • --csa-compress-rotary-base: 40000
  • --dsa-indexer-n-heads 64
  • --dsa-indexer-head-dim 128
  • --dsa-indexer-topk 512
• Sparse mode:
  • --experimental-attention-variant dsv4_hybrid
  • --csa-window-size 128
  • --csa-compress-ratios ([0,0,4]+[128,4]*20+[0])
  • --csa-compress-rotary-base: 40000
  • --dsa-indexer-n-heads 64
  • --dsa-indexer-head-dim 128
  • --dsa-indexer-topk 512
  • --dsa-indexer-use-sparse-loss

⚠️ 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 @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• 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

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

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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.

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

[cuichenx]

should we onboard the dev branch unit tests to this PR as well?

[cuichenx]

Is mHC meant to be included in this PR? I believe this arg is not yet on the main branch. The PR description also doesn't mention mHC

[cuichenx]

there are a few more CSA arguments not exposed here. is that intentional?

[FDecaYed]

Recording the replacement plan for the DeepSeek-V4 HybridModel recut: this historical PR is no longer an active merge path. The authoritative guide is DeepSeek-V4 HybridModel merge plan #5795.

The retained functionality is being reconstructed as:

• P6 / Explicit safe MLA RoPE primitives and consumers #5944 — explicit safe MLA RoPE primitives and consumers
• D1 — unfused SBHD compressed sparse attention
• D2 — SBHD DSv4 orchestration and indexer-loss scaling
• D4 — HybridModel DSv4 symbols, specs, and construction

GPTModel experimental specs and wiring are outside the frozen HybridModel-only scope. The original work remains linked from #5795 for provenance and credit. This PR will stay open as a draft reference, and the deepseekv4 label is being removed so it is not tracked as an active delivery PR.