Skip to content

Enable Nemotron 4 text inference CUDA graphs with Shortcut MoE and wide residuals - #7772

Open
susavlsh10 wants to merge 6 commits into
NVIDIA:mainfrom
susavlsh10:nemotron4-inference-cuda-graphs
Open

Enable Nemotron 4 text inference CUDA graphs with Shortcut MoE and wide residuals#7772
susavlsh10 wants to merge 6 commits into
NVIDIA:mainfrom
susavlsh10:nemotron4-inference-cuda-graphs

Conversation

@susavlsh10

@susavlsh10 susavlsh10 commented Oct 1, 2026 •
edited
Loading

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

What does this PR do?

Enable local, block-scoped CUDA graphs for Nemotron 4 text inference with GDP, Shortcut MoE, and wide residuals.

The change fixes GDP decode padding during graph replay and the FP32 wide-residual to BF16 RMSNorm precision boundary. It also restores the saved expert biases when loading MIMO checkpoints trained with global-batch quantilebalancing, so checkpoint routing is correct in eager and graph inference.

Regression tests cover configuration validation, GDP padding, RMSNorm dtype, checkpoint routing, and eager-versus-graph generation through a small Shortcut MoE model.

Dependencies

Depends on #7752 (text-only MIMO checkpoint loading), #6855 (wide residual MTP/MIMO support), and #7440 (wide residual Shortcut MoE support).

This draft currently includes a tested integration merge of the wide-residual work. We will refresh the dependency revisions and narrow the diff against main as those PRs progress.

Validation on tested revisions

  • NM4 nano derisk checkpoint matched accuracy eager vs graph
  • NM4 cuda graph inference gives ~6x speedups over eager baseline

⚠️ 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
  • If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
  • 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

copy-pr-bot Bot commented Oct 1, 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.

@susavlsh10
susavlsh10 marked this pull request as ready for review October 1, 2026 15:40
@susavlsh10
susavlsh10 requested review from a team as code owners October 1, 2026 15:40
mathemakitten
mathemakitten reviewed Oct 1, 2026
View reviewed changes
Comment thread megatron/core/tensor_parallel/inference_layers.py Outdated
Comment thread megatron/core/ssm/ssm_inference.py Outdated
mathemakitten
mathemakitten reviewed Oct 1, 2026
View reviewed changes
Comment thread megatron/core/transformer/transformer_layer.py Outdated
@mathemakitten mathemakitten mentioned this pull request Oct 2, 2026
Open
1 of 7 tasks
@susavlsh10
susavlsh10 force-pushed the nemotron4-inference-cuda-graphs branch from 92bdc80 to 753297e Compare October 3, 2026 00:39
santhnm2
santhnm2 reviewed Oct 3, 2026
View reviewed changes
Comment thread megatron/core/transformer/transformer_config.py
Comment thread megatron/training/arguments.py
Comment thread megatron/core/inference/text_generation_server/dynamic_text_gen_server/vlm_dynamic_inference.py Outdated
@lhb8125 lhb8125 mentioned this pull request Oct 5, 2026
Open
@sidsingh-nvidia

sidsingh-nvidia commented Oct 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

I suggest making the block scope default for inference. It is what we have been extensively using in Nemo-RL.

@susavlsh10

susavlsh10 commented Oct 5, 2026

Copy link
Copy Markdown
Author

I suggest making the block scope default for inference. It is what we have been extensively using in Nemo-RL.

@sidsingh-nvidia should we do this in this PR?

@Connor-XY Connor-XY mentioned this pull request Oct 6, 2026
Open
@santhnm2
santhnm2 removed request for a team October 6, 2026 16:10
@susavlsh10
Enable Nemotron 4 Shortcut MoE inference CUDA graphs
2cb2653 
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
@susavlsh10
Remove redundant inference RMSNorm dtype adaptation
5a6a958 
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
@susavlsh10
Align checkpoint routing test with upstream MIMO detection
cfa345b 
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
@susavlsh10
Handle SSM token padding independently of batch-invariant guards
7083a3b 
Always slice the decode partition and restore projection-local token padding. Require either the context or process-wide batch-invariant switch when token-only padding is present, and exercise both switches independently in the focused regression test.

Validation: 39 isolated CPU orchestration checks, Python syntax, Black, isort, Ruff, and git diff --check. CUDA capture/replay and distributed tests remain pending because the cluster is unavailable for GPU testing.
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
@susavlsh10
Prepare global-batch QB routing in shared inference configuration
4884d25 
Move the checkpoint router conversion out of MIMO text detection and into shared checkpoint and inference-config preparation. Retain the saved QB settings to support either initialization order, including prebuilt model configurations, without changing training routing.

Validation: 27 shared configuration checks; 40 focused regression cases on each of four ranks; fresh checkpoint-backed prefill/decode CUDA graph replay with exact eager token parity; changed-code syntax and formatting checks; git diff --check.
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
Susav Shrestha @susavlsh10
Remove QB inference workaround in favor of native support
1d5f7d4 
Signed-off-by: Susav Shrestha <susavlsh10@gmail.com>
@susavlsh10
susavlsh10 force-pushed the nemotron4-inference-cuda-graphs branch from 040471c to 1d5f7d4 Compare October 6, 2026 17:56

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Mellonta Mellonta Mellonta left review comments
@mathemakitten mathemakitten mathemakitten left review comments
@kwangjunahn kwangjunahn kwangjunahn left review comments
@santhnm2 santhnm2 santhnm2 approved these changes
@sidsingh-nvidia sidsingh-nvidia sidsingh-nvidia approved these changes

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@susavlsh10 @sidsingh-nvidia @santhnm2 @Mellonta @mathemakitten @kwangjunahn