[dev] [DeepSeek-v4] Context Parallel support

[kunlunl]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

Add CP support for DS v4

Depends on #5011 (Merged)

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

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

[kunlunl]

/claude strict-review

[claude]

Code Review Summary

PR: [dev] [DeepSeek-v4] Context Parallel support
Files reviewed: 15 (8003 additions, 116 deletions)
Findings: CRITICAL: 0, IMPORTANT: 3, SUGGESTION: 4

---

Overview

This PR adds context parallel (CP) support for the DeepSeek-v4 (DSv4) hybrid attention architecture. The implementation introduces:

1. CP partition modes — contiguous (default) and two_chunk (load-balanced for long sequences)
2. Boundary hidden exchange — P2P communication of window rows between CP ranks via _LeftBoundaryExchange autograd function
3. CP-aware compressor/indexer — All-gather of compressed KV, repacking from rank-major to sequence-major order, and CP-aware top-K indexer scoring
4. CuTeDSL layout kernels — Fused RoPE, KV packing, compressor input preparation, and attention index building for CP layouts
5. Data schedule integration — DSv4-specific CP row partitioning in get_cp_slice_for_thd

The architecture is well-structured: csa_cp_utils.py provides the MCore-facing API layer, csa_cp_layout_kernels.py contains the CuTe device code, and the existing csa.py / deepseek_v4_hybrid_attention.py files are extended with CP dispatch logic. Process groups are passed through via ProcessGroupCollection or explicit cp_group parameters, following the CLAUDE.md guidance.

Most Impactful Findings

1. Inconsistent getattr usage in data_schedule.py — The experimental_attention_variant attribute is accessed with getattr(config, ..., None) in one place but directly as config.experimental_attention_variant in another. The direct access will raise AttributeError if config doesn't have the attribute.

2. Missing padding-row exclusion in FusedIndexerSparseAttnFromTopkFunc — The new CP-specific autograd function lacks cu_seqlens_q_unpadded support for excluding CUDA-graph padding rows from the indexer loss. The existing non-CP counterpart (FusedIndexerSparseAttnFunc) has this support. This could cause slight training loss pollution when CP is used with CUDA graph capture and max-padded packing.

3. Dual indexer-loss implementations — The unfused CP path (_unfused_indexer_sparse_attn_from_topk) reimplements the target/predict KL loss computation independently from the shared helpers in dsa_kernels.py. Two independent implementations of the same math increase maintenance risk.

Risk Assessment

Medium risk. The core CP communication patterns (P2P boundary exchange, all-gather compressed KV) are correctly implemented with proper autograd backward support. The RoPE position mapping, KV packing, and index lowering kernels include appropriate validation. The new code is well-contained behind cp_size > 1 guards and does not alter single-rank behavior. The two-chunk partition mode adds complexity but is backed by comprehensive unit tests (1900+ lines of new test code).

The main risks are: (a) the padding-row exclusion gap for CUDA graph compatibility, and (b) the dual loss implementations potentially diverging over time.

[kunlunl]

/ok to test b83d57e

[kunlunl]

/ok to test d4832db

[kunlunl]

/ok to test e1198f6

[kunlunl]

indexer_top_k_wrapper compiles the kernel with bucketed_num_cols = next_pow2(num_cols), but scores_flat is still allocated with the original num_cols. When num_cols is not a power of two, the kernel may access beyond the allocated buffer. Padding scores_flat to the next power of two (using -inf) keeps the runtime tensor shape consistent with the compiled kernel and avoids potential out-of-bounds accesses.

--- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py
+++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py
@@ -599,6 +599,15 @@ def _indexer_topk_core(
 
     # ---------------- Shared: radix top-K + pad-to-topk -----------------
     topk_k = min(topk, sk)
+
+    # indexer_top_k_wrapper compiles the kernel with bucketed_num_cols=next_pow2(num_cols),
+    # but allocates buffer_torch with actual num_cols — causing out-of-bounds when num_cols
+    # is not a power of 2. Pad scores_flat to the next power of 2 to keep them consistent.
+    _ncols = scores_flat.shape[-1]
+    _ncols_p2 = 1 << (_ncols - 1).bit_length() if _ncols & (_ncols - 1) else _ncols
+    if _ncols_p2 != _ncols:
+        scores_flat = torch.nn.functional.pad(scores_flat, (0, _ncols_p2 - _ncols), value=float('-inf'))
+
     tk_result = _DSA.indexer_top_k_wrapper(
         scores_flat, seq_lens, top_k=topk_k, next_n=1, return_val=False

@xlsunstar Thanks for flagging this and for suggesting a fix. bucketed_num_cols is used as part of the compilation-cache key and the kernel specialization capacity. In my understanding its purpose is to let nearby runtime widths reuse the same compiled kernel and avoid unnecessary recompilation as num_cols changes; it is not intended to define the valid memory-access range.
The actual per-row access range is defined by seq_lens, and the kernel is expected to predicate its accesses using that runtime length. MCore sets it accordingly and guarantees seq_lens <= scores_flat.shape[-1] in both the THD and BSHD paths.
Have you observed an actual out-of-bounds access or correctness issue here? If so, it may indicate that the cuDNN top-k kernel is not fully honoring the seq_lens access-range contract and is instead using the bucketed compilation capacity as the runtime bound. In that case, we would prefer to fix the issue in cuDNN Frontend rather than pad the dense score tensor in MCore.

[chtruong814]

If possible, some of the dependency updates should be handled in pyproject toml and uv lock. I'd like @balasaajay and @ko3n1g to weigh in on bumping to the cudnn frontend commit.

[chtruong814]

I think flash_mla is already using the commit you want
https://github.com/NVIDIA/Megatron-LM/blob/main/pyproject.toml#L228

[chtruong814]

Oh this PR is going into dev branch. I would just cherry pick this change from main branch
https://github.com/NVIDIA/Megatron-LM/pull/5448/changes

[kunlunl]

Thanks for pointing this out. I mirrored the change from #5448 on the dev branch, pyproject.toml now tracks FlashMLA nv_dev, and uv.lock resolves it to b7643bd. I also removed the manual FlashMLA installation from the Dockerfile.

[chtruong814]

Please use the pyproject.toml and uv.lock to manage dependencies if possible

[kunlunl]

Done. FlashMLA and cuDNN Frontend are now declared in pyproject.toml, with the exact revisions and the CuTe DSL / DLPack transitive dependencies captured in uv.lock. The Dockerfile now installs this dependency chain through uv sync --locked; the separate manual installation has been removed.

[chtruong814]

@kunlunl Do we know when cudnn-frontend will include this in a release? Does this need to be included in next 26.08 Megatron-Bridge container release? Is the code backwards compatible to handle if cudnn frontend is not at this commit?

@balasaajay @ko3n1g any concerns?

[chtruong814]

@balasaajay @ko3n1g I didn't realize this was going into dev branch. I don't think much of a concern on the specific cudnn frontend used for dev branch, right? Main recommendation would still be to move what we can to pyproject toml and uv lock.

[kunlunl]

Thanks for checking. Since this PR targets the dev branch, I pinned the required cuDNN Frontend commit in pyproject.toml / uv.lock for the dev environment.

Older cuDNN Frontend versions remain usable for paths that do not require the new CP offset API. When ratio-4 fused CP is enabled, TransformerConfig now checks that the required wrappers expose q_causal_offsets and raises a clear initialization error if they do not, with the option to disable DSA kernel fusion.

[chtruong814]

Why do we need to reinstall? Can we handle this in pyproject toml and uv lock?

[chtruong814]

@balasaajay @ko3n1g btw, looks like we had to do this in MBridge container though I can't recall why.
NVIDIA-NeMo/Megatron-Bridge@9e06ad9

[kunlunl]

The explicit reinstall is no longer needed and has been removed. nvidia-cudnn-frontend[cutedsl] now resolves the required CuTe DSL 4.5.0 packages through uv.lock.

[kunlunl]

/ok to test 3599a37

[yuzhongw-nvidia]

LGTM! Thanks for your incredible work!

[kunlunl]

/ok to test 1543950

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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