[yuzhongw-nvidia]

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

What does this PR do ?

PR for dev: #5131

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

Summary

This PR avoids allocating dummy attention masks for TransformerEngine CUDA graph attention when attention masks are intentionally omitted by the dataloader.

When --no-create-attention-mask-in-dataloader is used, the eager forward path passes attention_mask=None. The TE CUDA graph path, however, was still creating synthetic all-zero masks for capture/replay with shape [micro_batch, 1, local_seq_len, global_seq_len]. With context parallelism, scaling seq_len and CP together keeps local sequence length constant, but the dummy mask still scaled with global seq_len, causing unnecessary CUDA graph memory growth.

Changes

• Plumb create_attention_mask_in_dataloader into TransformerConfig.
• For TE CUDA graph attention with create_attention_mask_in_dataloader=False, omit attention_mask from both static capture inputs and replay kwargs instead of allocating a dummy mask.
• Preserve the previous mask behavior when dataloader-created masks are enabled.
• Add warnings for attention mask types where omitting the mask may be unsafe.
• Update unit coverage for TE CUDA graph replay/static input behavior.

Validation

Validated on a Blackwell platform using a 4-layer Qwen3-30B debug recipe with partial CUDA graph enabled.

The key comparison is the same sequence-length / CP cases before vs. after the fix:

Case	Before fix peak allocated	Before fix peak reserved	Before fix device memory	After fix peak allocated	After fix peak reserved	After fix device memory
16k-CP1	29,258 MB	32,702 MB	41,219 MB	29,258 MB	31,942 MB	40,459 MB
32k-CP2	29,263 MB	34,666 MB	44,265 MB	29,263 MB	32,424 MB	42,023 MB
64k-CP4	31,128 MB	36,612 MB	46,601 MB	29,264 MB	32,626 MB	42,621 MB

For the long-sequence cases, the fix removes the CUDA-graph-specific dummy-mask memory:

Case	Peak allocated reduction	Peak reserved reduction	Device memory reduction
16k-CP1	0 MB	760 MB	760 MB
32k-CP2	0 MB	2,242 MB	2,242 MB
64k-CP4	1,864 MB	3,986 MB	3,980 MB

The growth from 16k-CP1 -> 64k-CP4 is also reduced:

• Before fix: device memory increased by 5,382 MB.
• After fix: device memory increased by 2,162 MB.

The remaining delta matches the no-CG / MoE-only CP baseline. This confirms that the attention CUDA graph path was retaining synthetic [local_seq_len, global_seq_len] attention masks, and that omitting those masks when --no-create-attention-mask-in-dataloader is set removes the long-sequence CP memory regression.

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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.