Skip to content

[PyTorch] Add torch.compile support for fused MoE routing #3691

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pggPL wants to merge 14 commits into
base: main
Choose a base branch
from
Open

[PyTorch] Add torch.compile support for fused MoE routing #3691

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
b335631
[PyTorch] Add torch.compile support for fused MoE routing
pggPL Oct 14, 2026
7f30b44
Support compiled dense router indices without overwriting buffers
pggPL Oct 14, 2026
bd8328c
Support compiled quantile balancing with functional histogram outputs
pggPL Oct 14, 2026
f5c02da
Reuse router numerical tests for selected torch.compile configurations
pggPL Oct 15, 2026
e0aa561
Compile router buffer mutations through autograd.Function
pggPL Oct 15, 2026
9f24dfc
Skip QB bounds value validation under torch.compile
pggPL Oct 15, 2026
a3a6a10
Simplify compiled router inference context handling
pggPL Oct 15, 2026
ea3275c
Remove unrelated router test formatting change
pggPL Oct 15, 2026
41219e9
Simplify fused router tests and helper names
pggPL Oct 15, 2026
80b7920
Remove dedicated compiled QB bounds validation tests
pggPL Oct 15, 2026
b8d8eec
Document router buffer constraints beside their parameters
pggPL Oct 15, 2026
9282a4b
Use setup_context for compiled fused routing
pggPL Oct 15, 2026
eb9e922
Check that dynamic router inputs do not trigger recompilation
pggPL Oct 16, 2026
c2fb139
Check fused router backward custom operators with opcheck
pggPL Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Loading