Skip to content

JIT for MoE router and preprocess - #1919

Merged
Victarry merged 11 commits into
NVIDIA:mainfrom
yaox12:xiny/jit_for_moe_main
Nov 19, 2025
Merged

JIT for MoE router and preprocess#1919
Victarry merged 11 commits into
NVIDIA:mainfrom
yaox12:xiny/jit_for_moe_main

Conversation

@yaox12

@yaox12 yaox12 commented Oct 24, 2025

Copy link
Copy Markdown
Member
  • Enable JIT (torch.compile) for expert bias and MoE preprocess.
  • Let force load balancing compatible with CUDA Graph.

Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12
yaox12 requested review from a team as code owners October 24, 2025 07:37
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2025

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.

@yaox12

yaox12 commented Oct 24, 2025

Copy link
Copy Markdown
Member Author

/ok to test b3283b9

Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
@yaox12 yaox12 added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Nov 5, 2025
@yanring yanring added this to the Core 0.16 milestone Nov 14, 2025
@yanring
yanring enabled auto-merge November 17, 2025 03:51
@yanring

yanring commented Nov 17, 2025

Copy link
Copy Markdown
Contributor

/ok to test 04b3623

@yanring
yanring added this pull request to the merge queue Nov 18, 2025
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Nov 18, 2025
@Victarry
Victarry added this pull request to the merge queue Nov 19, 2025
Merged via the queue into NVIDIA:main with commit efdc681 Nov 19, 2025
45 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. module: moe

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants