Prevent FlashInfer sampling from running with CUDA graphs by santhnm2 · Pull Request #5791 · NVIDIA/Megatron-LM

santhnm2 · 2026-07-13T22:42:27Z

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

What does this PR do?

Prevents FlashInfer sampling from running with CUDA graphs enabled due to a known issue regarding the RNG state not being CUDA-graph safe.
See here for more details: https://www.linkedin.com/pulse/pinned-rng-drifting-crash-from-cuda-graph-chenyang-zhao-csuac/

Also enables more fine-grained dispatching across sampling kernels depending on the batch composition (i.e., different kernels for only top-k, only top-p, and mixed batches).

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

When your PR is ready, click Ready for Review.
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.

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

copy-pr-bot · 2026-07-13T22:42:30Z

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 · 2026-07-13T22:44:00Z

/ok to test 50a7417

kvareddy · 2026-07-14T15:26:09Z

@tdene can you please take a look at this MR?

sidsingh-nvidia · 2026-07-14T15:37:55Z

Can we make flashinfer the default backend (just like vllm) and gracefully fallback to torch if it is not installed?

…ling

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-07-14T20:58:28Z

/ok to test f1c60b5

santhnm2 · 2026-07-14T20:58:49Z

Can we make flashinfer the default backend (just like vllm) and gracefully fallback to torch if it is not installed?

Done

santhnm2 · 2026-07-20T16:58:19Z

/ok to test d84e7b6

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-07-20T21:07:44Z

/ok to test 4fb14c2

svcnvidia-nemo-ci · 2026-07-20T22:38:02Z

🔄 Merge queue validation started!

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

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-07-21T06:42:34Z

/ok to test 519baaa

svcnvidia-nemo-ci · 2026-07-21T20:10:40Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-21T23:03:51Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-22T04:15:26Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-22T09:50:07Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-22T15:50:21Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-23T02:20:48Z

🔄 Merge queue validation started!

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

Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

sidsingh-nvidia · 2026-07-23T08:15:07Z

Pushed a commit to fix the failing test harness.

sidsingh-nvidia · 2026-07-23T20:48:16Z

/ok to test b571fcd

svcnvidia-nemo-ci · 2026-07-24T23:10:36Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-07-25T00:27:40Z

🔄 Merge queue validation started!

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

Prevent FlashInfer sampling from running with CUDA graphs

50a7417

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

copy-pr-bot Bot temporarily deployed to public July 13, 2026 22:44 Inactive

copy-pr-bot Bot temporarily deployed to test July 13, 2026 22:44 Inactive

copy-pr-bot Bot temporarily deployed to public July 13, 2026 22:47 Inactive

copy-pr-bot Bot temporarily deployed to public July 13, 2026 22:48 Inactive

copy-pr-bot Bot temporarily deployed to public July 13, 2026 23:05 Inactive

santhnm2 marked this pull request as ready for review July 14, 2026 15:22

santhnm2 requested review from a team as code owners July 14, 2026 15:22

svcnvidia-nemo-ci added the complexity: medium label Jul 14, 2026

kvareddy requested a review from tdene July 14, 2026 15:25

sidsingh-nvidia reviewed Jul 14, 2026

View reviewed changes

Comment thread megatron/core/inference/sampling/flashinfer_sampling.py Outdated

shanmugamr1992 approved these changes Jul 14, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Jul 14, 2026

santhnm2 added 2 commits July 14, 2026 13:44

Merge remote-tracking branch 'upstream/main' into fix_flashinfer_samp…

680d215

…ling

Default to flashinfer sampling, remove torch.multinomial

f1c60b5

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 requested review from a team as code owners July 14, 2026 20:58

svcnvidia-nemo-ci removed the Final Review PR is in the "final review" stage label Jul 14, 2026

copy-pr-bot Bot temporarily deployed to public July 14, 2026 20:59 Inactive

Merge branch 'main' into fix_flashinfer_sampling

d84e7b6

santhnm2 enabled auto-merge July 20, 2026 16:58

copy-pr-bot Bot temporarily deployed to public July 20, 2026 16:58 Inactive

copy-pr-bot Bot temporarily deployed to test July 20, 2026 16:59 Inactive

copy-pr-bot Bot temporarily deployed to public July 20, 2026 17:02 Inactive

santhnm2 added 3 commits July 20, 2026 14:06

Merge with main

bbae0a4

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Linting

96fb397

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Merge branch 'main' into fix_flashinfer_sampling

4fb14c2

santhnm2 added 2 commits July 20, 2026 23:41

Update functional test golden values

f2fd129

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Merge branch 'main' into fix_flashinfer_sampling

519baaa

fix inference test

1f2663c

Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>

ping PR to refresh

b571fcd

thomasdhc approved these changes Jul 24, 2026

View reviewed changes

Uh oh!

Conversation

santhnm2 commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jul 13, 2026

Uh oh!

santhnm2 commented Jul 13, 2026

Uh oh!

kvareddy commented Jul 14, 2026

Uh oh!

Uh oh!

sidsingh-nvidia commented Jul 14, 2026

Uh oh!

santhnm2 commented Jul 14, 2026

Uh oh!

santhnm2 commented Jul 14, 2026

Uh oh!

santhnm2 commented Jul 20, 2026

Uh oh!

santhnm2 commented Jul 20, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 20, 2026

Uh oh!

santhnm2 commented Jul 21, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 21, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 21, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 22, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 22, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 22, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 23, 2026

Uh oh!

sidsingh-nvidia commented Jul 23, 2026

Uh oh!

sidsingh-nvidia commented Jul 23, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 24, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

santhnm2 commented Jul 13, 2026 •

edited

Loading