Skip to content

Enable CUDA graphs for MTP inference#4260

Merged
santhnm2 merged 148 commits into
NVIDIA:mainfrom
santhnm2:mtp_perf
Apr 26, 2026
Merged

Enable CUDA graphs for MTP inference#4260
santhnm2 merged 148 commits into
NVIDIA:mainfrom
santhnm2:mtp_perf

Conversation

@santhnm2

@santhnm2 santhnm2 commented Apr 11, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Enables CUDA graphs for MTP inference.

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 the @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.

santhnm2 added 22 commits April 9, 2026 21:06
@santhnm2
Add nvtx ranges for mtp
bde9685 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
perf benchmarking
7587f2f 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Remove simpy dependency
b8bdd12 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
cuda graphs for mtp
e948bc9 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
mtp cuda graphs
6c60f34 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
fix cuda graphs
90ffe8a 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
cuda graph fix
8ce0a4b 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Fix dummy_position_ids dtype
84129e7 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Try to compile more
9bfa0fc 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
remove extra ep sync
cbc1270 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Add nvtx ranges
7687743 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Remove cpu sync
407e12d 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Avoid h2d syncs
e54458c 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Linting
70afe4c 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
compile rewind_kv_cache
5e84f85 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Linting
e2020b2 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Fix verify_speculative_tokens
56b93ff 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
More fixes
98a32ef 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Add torch.compile back
3fefb93 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Linting
96668d6 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Try to fix tests
9a95577 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Fix test
74edd58 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Apr 11, 2026

Copy link
Copy Markdown

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.

santhnm2 added 4 commits April 11, 2026 00:53
@santhnm2
More test fixes
b6466d2 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Bug fix
6941afb 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Linting
94452d5 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Merge remote-tracking branch 'upstream/main' into mtp_perf
cd981ea
@santhnm2

santhnm2 commented Apr 11, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test cd981ea

@svcnvidia-nemo-ci svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 11, 2026
santhnm2 added 4 commits April 24, 2026 12:24
@santhnm2
Add garbage collection to test_moe_inference.py
9f64401 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Linting
f193a01 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2
Merge remote-tracking branch 'upstream/main' into mtp_perf
5266808
@santhnm2
More cuda graph deletion
c6ad4f3 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2

santhnm2 commented Apr 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 5266808

@santhnm2
Fix copyright
2cae20f 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2

santhnm2 commented Apr 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2cae20f

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Apr 24, 2026
@santhnm2
Revert unnecessary cuda graph changes
f537865 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@santhnm2

santhnm2 commented Apr 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f537865

@santhnm2
Merge with main
84b1ef5 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Apr 26, 2026
@santhnm2

santhnm2 commented Apr 26, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 84b1ef5

@santhnm2 santhnm2 enabled auto-merge April 26, 2026 17:10
@santhnm2 santhnm2 added this pull request to the merge queue Apr 26, 2026
@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Apr 26, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 03f4111 Apr 26, 2026
348 of 350 checks passed
@tdene tdene mentioned this pull request Apr 28, 2026
Merged
5 tasks
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
@santhnm2
Enable CUDA graphs for MTP inference (NVIDIA#4260)
d1e9ad8 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@sbhavani sbhavani mentioned this pull request May 26, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jaredcasper jaredcasper jaredcasper approved these changes

@kvareddy kvareddy kvareddy approved these changes

@mathemakitten mathemakitten mathemakitten approved these changes

@lmcafee-nvidia lmcafee-nvidia lmcafee-nvidia approved these changes

@shanmugamr1992 shanmugamr1992 shanmugamr1992 approved these changes

@sidsingh-nvidia sidsingh-nvidia sidsingh-nvidia approved these changes

@claude claude[bot] claude[bot] approved these changes

Assignees

No one assigned

Labels

Approved All necessary approvals have been made complexity: high Run functional tests

Projects

None yet

Milestone

Core 0.16

Development

Successfully merging this pull request may close these issues.

8 participants

@santhnm2 @svcnvidia-nemo-ci @jaredcasper @kvareddy @mathemakitten @lmcafee-nvidia @shanmugamr1992 @sidsingh-nvidia