MInf: bound the FlashInfer MoE expert kernels by token count by tdene · Pull Request #8367 · NVIDIA/Megatron-LM

tdene · 2026-10-20T16:24:07Z

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

What does this PR do?

This makes the FlashInfer MoE kernels only do work proportional to the live tokens on a step, instead of proportional to the overall buffer. This in the same spirit of changes started by #4258.

As a result, my testing shows that a decode step on Nano 3.5 drops from 25ms to 9ms with a 65,536-row buffer (EP4), and from 76ms to 12ms with a 262,144-row buffer (EP16), with no change at all in outputs.

Problem

The FlashInfer MoE kernels currently assume that the token count can be provided as a host-side parameter. This is false. The host only provides this parameter at CUDA graph record, and it has no way of knowing what the token count will ever be, because the token count is only known after the all-to-all expert collectives run on every layer.

A secondary problem is that the FlashInfer MoE kernels currently downcast their accumulated results from FP32 to BF16 and only return BF16, forcing the user (Megatron Inference) to manually upcast the BF16 out back to FP32. This is a loss of precision and slowdown with no benefit.

Solution

Monkey-patch the FlashInfer kernels until they are fixed at a future date.

Perf numbers

One GB300 node, Nemotron 3.5 Nano, TP 2 EP 2 DP 2, 256 requests of 128-token prompts, 1,024 greedy tokens, three
trials per run averaged out:

decode step timing	65,536 rows	131,072 rows	262,144 rows
bf16, `main` with FlashInfer	23.6ms	38.0ms	67.9ms
bf16, this PR with current FlashInfer	9.8ms	10.8ms	13.3ms
bf16, this PR with a future, fixed, FlashInfer	9.3ms	10.3ms	11.9ms

decode step timing	65,536 rows	131,072 rows	262,144 rows
MXFP8, `main` with FlashInfer	25.3ms	42.5ms	76.2ms
MXFP8, this PR with current FlashInfer	9.2ms	10.6ms	12.1ms
MXFP8, this PR with a future, fixed, FlashInfer	8.9ms	9.8ms	11.8ms

Please see the Slack thread for a wide list of perf numbers. More are available upon request.

Related PRs

flashinfer-ai/flashinfer#6783 flashinfer-ai/flashinfer#6784 flashinfer-ai/flashinfer#6785

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"

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.

copy-pr-bot · 2026-10-20T16:24:11Z

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.

mathemakitten

Right now the temporary parts of this PR are mixed with the permanent ones, and unwinding them after the FlashInfer PRs land would be another review cycle. Is it possible to restructure this in a way to minimize the changes and make unwinding a matter of just deleting a temp module?

I would also like to understand the speedup in decode vs main with inference_flashinfer_mxfp8_token_capacity set (for MXFP8). The comparison to make is not actually main with unbounded rows but main with bounded rows (which is only available in MXPF8 right now, but there is no real reason it can't be made available to BF16). Can we run that comparison?

tdene · 2026-10-28T19:39:31Z

Right now the temporary parts of this PR are mixed with the permanent ones, and unwinding them after the FlashInfer PRs land would be another review cycle. Is it possible to restructure this in a way to minimize the changes and make unwinding a matter of just deleting a temp module?

Done; how does this look?

tdene · 2026-10-28T22:53:57Z

I would also like to understand the speedup in decode vs main with inference_flashinfer_mxfp8_token_capacity set (for MXFP8). The comparison to make is not actually main with unbounded rows but main with bounded rows (which is only available in MXPF8 right now, but there is no real reason it can't be made available to BF16). Can we run that comparison?

MXFP8 decode, 256 requests, 65,536 rows	tok/s per replica	ms per step
main, no bound	10,127	25.3
main with `inference_flashinfer_mxfp8_token_capacity=4096`	24,285	10.5
this PR, current FlashInfer and no bound	27,687	9.2
this PR with FlashInfer 6783, 6784	28,612	8.9
this PR with FlashInfer 6783, 6784, 6785	30,425	8.4

There's a few problems I see with the bound solution. One problem is that it only runs on decode steps by design, which means that when we have frequent prefills we do not see a benefit. Since EP > DP, if any of the DP replicas are in prefill, none of them will use the bound.

I'm working to get that kind of data out, just hitting various issues with main.

balasaajay · 2026-10-29T03:05:00Z

/ok to test b21432e

mathemakitten

Some cleanup items, but otherwise the restructure is looking great.

mathemakitten · 2026-11-02T16:12:59Z

+
+
+@requires_cuda_triton
+def test_routed_emulation_finalizes_the_live_rows(monkeypatch):


These emulation tests stub the kernel, so they check the Triton finalize/combine arithmetic but not the assumption underneath it: that do_finalize=False returns unweighted FC2 rows (routed) and that enable_alltoall=True writes exactly the locally-routed rows (cutlass). A double- or missing-weight bug would pass these and the replay tests.

Could we add a Blackwell-only test per kernel that runs the real FlashInfer kernel both ways on the same inputs and compares rows [0, num_active_tokens): the shim's emulated path vs. the kernel's own finalize (do_finalize=True for routed, default enable_alltoall=False for cutlass) over all rows?

Yeah, thanks for pointing this out. Should be fixed now.

mathemakitten · 2026-11-02T16:17:23Z

+
+# UNWIND(flashinfer num_active_tokens): once flashinfer-ai/flashinfer#6783 ships,
+# import cutlass_fused_moe from flashinfer.fused_moe instead.
+from megatron.core.inference.moe.flashinfer_active_tokens import cutlass_fused_moe


Can we guard this in a try block like the rest of the FlashInfer imports? Lots of containers (incl likely training containers) are still missing FlashInfer and we cannot break on importing experts.py.

The cutlass_fused_moe file already has a try/catch around FlashInfer importing; experts.py is safe to import without FlashInfer.

mathemakitten · 2026-11-02T16:22:33Z

+    do_finalize: bool = True,
+    output: torch.Tensor | None = None,
+    num_active_tokens: torch.Tensor | None = None,
+    **kwargs,


We do not accept **kwargs in /core, please pass explicit args.

My bad! Fixed

mathemakitten · 2026-11-02T16:24:18Z

+                torch.float32,
+                quant_scales=None,
+                activation_type=self._flashinfer_activation_type,
+                **(


We do not accept **kwargs in core, can we pass args unconditionally? uv.lock pins FlashInfer 0.7.2.post1 so this shouldn't be a problem, though you might have to update pyproject.toml.

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

kvareddy · 2026-11-03T06:28:16Z

@sidsingh-nvidia @wdykas can you please take a look at this MR?

tdene marked this pull request as ready for review October 20, 2026 16:24

tdene requested review from a team as code owners October 20, 2026 16:24

copy-pr-bot Bot temporarily deployed to public October 20, 2026 16:28 Inactive

github-actions Bot added the complexity: high label Oct 20, 2026

copy-pr-bot Bot temporarily deployed to test October 20, 2026 22:03 Inactive

copy-pr-bot Bot temporarily deployed to public October 20, 2026 22:06 Inactive

yaox12 approved these changes Oct 21, 2026

View reviewed changes

mathemakitten reviewed Oct 28, 2026

View reviewed changes

tdene force-pushed the tde/flashinfer_moe_active_tokens branch from d6bdd10 to 9dd8b1e Compare October 28, 2026 18:24

copy-pr-bot Bot had a problem deploying to test October 28, 2026 18:28 Error

copy-pr-bot Bot temporarily deployed to public October 28, 2026 18:32 Inactive

copy-pr-bot Bot had a problem deploying to test October 28, 2026 19:44 Error

copy-pr-bot Bot temporarily deployed to public October 28, 2026 19:47 Inactive

copy-pr-bot Bot temporarily deployed to test October 29, 2026 03:06 Inactive

copy-pr-bot Bot temporarily deployed to public October 29, 2026 04:28 Inactive

kahyunnam mentioned this pull request Oct 29, 2026

[RL] CUTLASS and TRT-LLM MoE: FP32 finalize output, device-side active token count under CUDA graphs, static GEMM configs (#6783, #6784, #6785) flashinfer-ai/flashinfer#7922

Open

lhb8125 mentioned this pull request Nov 1, 2026

Megatron-LM Monthly Activity Report — October 2026 #8811

Open

mathemakitten reviewed Nov 2, 2026

View reviewed changes

tdene and others added 4 commits November 2, 2026 11:03

MInf: bound the FlashInfer MoE expert kernels

bb7c00b

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

lint

0bcb316

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Refactor

7ab6f81

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

Address reviewer comments

16dc3f7

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

tdene force-pushed the tde/flashinfer_moe_active_tokens branch from b21432e to 16dc3f7 Compare November 2, 2026 18:20

copy-pr-bot Bot deployed to test November 2, 2026 18:24 Active

copy-pr-bot Bot temporarily deployed to public November 2, 2026 18:26 Inactive

mathemakitten approved these changes Nov 2, 2026

View reviewed changes

kvareddy requested review from sidsingh-nvidia and wdykas November 3, 2026 06:27



		@requires_cuda_triton
		def test_routed_emulation_finalizes_the_live_rows(monkeypatch):

Conversation

tdene commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Problem

Solution

Perf numbers

Related PRs

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 Oct 20, 2026

Uh oh!

mathemakitten left a comment

Choose a reason for hiding this comment

Uh oh!

tdene commented Oct 28, 2026

Uh oh!

tdene commented Oct 28, 2026

Uh oh!

balasaajay commented Oct 29, 2026

Uh oh!

mathemakitten left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kvareddy commented Nov 3, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tdene commented Oct 20, 2026 •

edited

Loading