Add packed THD context parallelism for DSv4

[FDecaYed]

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

What does this PR do?

DSv4 attention currently rejects CP > 1. This adds static context parallelism for packed THD inputs with attention_cp_layout="contiguous" and TP=1, covering window attention, ratio-4 CSA, ratio-128 HCA, and MTP. Packed inputs require dsa_kernel_backend="cudnn"; the existing SBHD native backend remains available.

• Exchange boundary hidden states, preserve document-relative RoPE, and assign compressed groups to their owning CP ranks. Handle packed padding, empty documents, and groups crossing rank boundaries.
• Connect packed FlashMLA/cuDNN attention and sparse/dense indexer loss, with asynchronous compressed-KV/indexer-K gathers, gradient reduce-scatter overlap, and KV reconstruction for selective recompute.
• Add contiguous MTP token/mask rolls, batch each phase's field exchanges, and return HSM boundary gradients to their original ranks.
• Add CP2/CP4 output and gradient comparisons against per-document native SBHD references, boundary/roll coverage, and kernel determinism registrations.

The implementation uses main's existing packing, process-group, RoPE, and static-spec interfaces. It preserves one dense indexer scoring call per CP-local query shard. Compact BF16/MXFP8 indexers and dynamic CP remain separate follow-ups; dependency versions are unchanged. This PR targets main independently of #6403.

Related work

Related to #5795 and #6404. Static CP behavior follows the existing dev CSA implementation, adapted to main's interfaces and pinned cuDNN Frontend.

Validation

• Passed 96 before/after CPU custom-autograd comparisons for sparse/dense loss, padding, CP1/2/4, and KV reconstruction, plus four pure-attention comparisons. External CUDA kernels and collectives were explicit substitutes; these checks verify adapter and gradient-routing equivalence.
• Passed CPU checks of the per-document reference outputs/gradients, padding, and global-token normalization for dense/sparse indexer loss.
• Passed AST, Black/isort, Pylint, Ruff, git diff --check, and kernel determinism coverage checks. Local mypy still reports dependency/type-configuration errors.
• Real GPU unit tests, full training, peak memory, and throughput have not been run locally. CPU checks do not qualify the external GPU kernels. Existing SM90 dense-indexer-loss and dense empty-compressed-segment restrictions remain.

GPU validation targets:

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_packed_cp.py \
  tests/unit_tests/transformer/test_contiguous_sequence_roll.py \
  tests/unit_tests/determinism/kernels/test_csa_packed_kernels.py \
  tests/unit_tests/determinism/kernels/test_fused_triton_kernels.py

Pre-checks

• Added relevant unit tests.
• Added relevant functional tests.
• Added/updated kernel determinism tests and registered them in the manifest.
• Completed typing validation.
• Ran tools/autoformat.sh; blocking formatter/lint checks passed.

[copy-pr-bot]

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.

[FDecaYed]

/ok to test e17057c

[FDecaYed]

/ok to test ef78927

[FDecaYed]

/ok to test 8e11d6b

[Connor-XY]

Hi @FDecaYed, a few notes from stacking #7705/#7708 on this:

1. Under torch.compile, bucketize(rows, cu[1:]) ignores the slice offset in our PyTorch build, so rows land one document to the right. Fix compiled packed sequence boundary lookup #7705 has the fix plus regressions. Could you adopt it here, or merge Fix compiled packed sequence boundary lookup #7705 right after this?
2. Rebasing onto main would bring in cuDNN frontend 1.29.0, which Make DSv4 attention and indexer backward deterministic #7708's deterministic backward needs.
3. The packed path ignores deterministic_mode until Make DSv4 attention and indexer backward deterministic #7708 lands. Should it raise in the meantime?

[Phlip79]

Please add docstring

[FDecaYed]

Documented cp_layout, fill_value and the contiguous-only shift/fill contract in 0f2ced4.

[Phlip79]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Summary

This is a strict review of the DSv4 packed-THD static context-parallel PR. I read all 17 changed files in the head tree. I also followed the paths that cross components:

• the hybrid attention, and the CSA compressor and indexer;
• the packed layout and the async all-gather/reduce-scatter (AG/RS) helpers;
• the MTP rolls, including the zigzag, CP1 and HSM paths;
• the config validation.

Findings: 0 CRITICAL · 0 IMPORTANT · 7 SUGGESTION (all inline)

• Four new args or branches that never take effect. The "every new arg has a real use path" check fails for each:
  • fill_value=True in the zigzag MTP padding_mask roll does nothing.
  • compressed_base and compressed_is_sequence_major in build_attention_indices are never passed.
  • The packed branches of _apply_unfused_rope cannot be reached.
  • The second operand of the guard in compute_cp_indexer_topk is always false.
• Missing docs: the roll_tensor docstring does not cover cp_layout or fill_value. This matches Phlip79's open request.
• Lost diagnostic: fused_csa_teacher_lse now raises a generic ValueError. The base raised a RuntimeError that included the specific unsupported reason.
• Untested path: every packed CP test fixes apply_rope_fusion=True, so none runs the new unfused-RoPE code.

Risk: moderate. This is a large new CP surface. Neither the author nor I could run its GPU kernels or collectives. Reading the source, I found no correctness, process-group, CPU-GPU sync, CUDA-graph or backward-compatibility defect. Since nothing above is significant, the verdict is approve. Please check the suggestions before merging.

What I checked by reading the source (not by running it):

• Compressed-group ownership: each compressed group belongs to the rank that holds its last token. The 8-row halo always contains the ratio-4 overlap predecessor, so is_first = group_id == 0 is correct.
• Row mapping: seq_to_rank_row maps correctly into the capacity-padded, rank-major gathered buffer. Halo copies never enter the global KV domain.
• Deferred reduce-scatter edges: they return each compressed-KV and indexer-K gradient exactly once. The gathered buffers are detached on the loss path.
• Loss scaling: real_row_scale matches loss_divisor for the sparse loss, the dense loss and per-token loss.
• Roll backward: _RollLeft and _LeftBoundaryExchange send boundary gradients back to the rank that owns the token.
• Process groups: every P2P call and collective passes the explicit CP group.
• Host syncs: none were added. The code uses torch._assert_async, and its boolean-mask scalar fills dispatch to masked_fill.
• Unchanged behavior: zigzag and CP1 MTP behave as before. For CP1, real and physical boundaries are the same tensor.

Not assessed: the PR discussion reports that bucketize(rows, cu[1:]) ignores the slice offset under torch.compile on some PyTorch builds. I cannot confirm this from source. _compiled_compactor and _compiled_attention_indices both use this pattern, so the 4-GPU test_dsv4_packed_cp.py run is what will show whether it holds here.

Review generated by claude

[FDecaYed]

@Connor-XY Adopted #7705's boundary fix and metadata regressions in 0f2ced4, with co-author credit. Its scoped changes are now included here. Packed deterministic_mode now raises explicitly pending #7921, with #7708 as the shared/SBHD prerequisite. I kept the existing dependency pins in this PR; the deterministic follow-up can use main's newer runtime. New CI is running.

[nemo-automation-bot]

🔄 Merge queue validation started!

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