Moe support #3811

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

Jump to bottom

Open

cehongwang wants to merge 8 commits into main from moe-support

+264 −97

Collaborator

cehongwang commented Sep 9, 2025

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

meta-cla bot added the cla signed label

github-actions bot added component: tests component: lowering component: conversion component: converters component: api [Python] component: dynamo labels

github-actions bot requested a review from narendasan

September 9, 2025 19:29

cehongwang force-pushed the moe-support branch 2 times, most recently from fa53986 to 6ea89ae Compare

September 17, 2025 22:06

github-actions bot added the component: runtime label

peri044 reviewed

View reviewed changes

py/torch_tensorrt/dynamo/conversion/impl/select.py

Collaborator

peri044 Sep 24, 2025

Could you add detailed comments demonstrating what it does now that you've gone through the entire converter ? that would be helpful

tools/llm/run_llm.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py Outdated Show resolved Hide resolved

cehongwang force-pushed the moe-support branch from 9e7ca5d to c286767 Compare

September 24, 2025 23:33

github-actions bot removed the component: runtime label

narendasan reviewed

View reviewed changes

py/torch_tensorrt/dynamo/lowering/passes/remove_num_users_is_0_nodes.py

    
                          gm.graph.erase_node(node)

                          gm = clean_up_graph_after_modifications(gm)

                  gm = clean_up_graph_after_modifications(gm)

Collaborator

narendasan Oct 7, 2025

Is this just because we only need one cleanup?

Collaborator Author

cehongwang Oct 7, 2025

Yeah or otherwise it could take forever

peri044 reviewed

View reviewed changes

tests/py/dynamo/conversion/test_index_put_aten.py Show resolved Hide resolved

tests/py/dynamo/conversion/test_index_put_aten.py Show resolved Hide resolved

narendasan approved these changes

View reviewed changes

Collaborator

narendasan left a comment

LGTM but please also get @peri044's stamp

peri044 approved these changes

View reviewed changes

Collaborator

peri044 commented Oct 8, 2025

@cehongwang the linux tests are failing. Do you know the rootcause ?

cehongwang added 2 commits

October 8, 2025 23:11


          Enabled Qwen MoE with 1 layer. Rewrote index_put converter

5154f4d


          fixed the perf issue in the lowering pass

544c038

cehongwang added 4 commits

October 8, 2025 23:11


          Optimized index converter

9aaa4ca


          Fixed a typo in the converter. Covered the discontinuous tests

d51370e


          Supported bool mask indicies

06e3a04


          Fixed the comments

049b810

cehongwang force-pushed the moe-support branch from 92f62f8 to 049b810 Compare

October 8, 2025 23:11

Collaborator Author

cehongwang commented Oct 8, 2025

@cehongwang the linux tests are failing. Do you know the rootcause ?

fixed


          Fixed the CI

6cf31cf

cehongwang force-pushed the moe-support branch from 015bf19 to 6cf31cf Compare

October 9, 2025 18:48


          Merge branch 'main' into moe-support

30471bb

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

Labels

cla signed component: api [Python] component: conversion component: converters component: dynamo component: lowering component: tests