Skip to content

[TPU] Avoid Triton Import #15589

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.

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

Merged

Conversation

robertgshaw2-redhat
Copy link
Collaborator

@robertgshaw2-redhat robertgshaw2-redhat commented Mar 27, 2025

SUMMARY:

  • delay triton import until needed
  • breaks tpu backend

Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Copy link

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

@robertgshaw2-redhat robertgshaw2-redhat changed the title fix triton import [TPU] Avoid Triton Import Mar 27, 2025
Copy link
Collaborator

@houseroad houseroad left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good.

Btw, shall we consider to have a common interface of layers, and implementations for different backend are put in the different files.

@robertgshaw2-redhat
Copy link
Collaborator Author

robertgshaw2-redhat commented Mar 27, 2025

Looks good.

Btw, shall we consider to have a common interface of layers, and implementations for different backend are put in the different files.

yes, I can catch you up @houseroad on what our plan is offline in a separate thread. We have been planning a refactor led by Lucas but its just been hard to prioritize it with all the models coming out. If your team has capacity to help with it, we can collaborate

@robertgshaw2-redhat robertgshaw2-redhat enabled auto-merge (squash) March 27, 2025 02:41
@github-actions github-actions bot added the ready ONLY add when PR is ready to merge/full CI is needed label Mar 27, 2025
Copy link
Member

@mgoin mgoin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, unfortunately

@mgoin mgoin added bug Something isn't working tpu Related to Google TPUs labels Mar 27, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Copy link
Collaborator

@yaochengji yaochengji left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

@wangxiyuan
Copy link
Contributor

+1, break other backend as well.

@MengqingCao
Copy link
Contributor

Thanks for this fix! BTW, I'm doing smilar things in pr #15099.
Do you think adding a defensive check before importing trion, and adding dummy decorators for triton is a good way to avoid triton breaks like this? Please let me know your suggestions, thx!

@robertgshaw2-redhat robertgshaw2-redhat merged commit e1e0fd7 into vllm-project:main Mar 27, 2025
35 checks passed
@robertgshaw2-redhat
Copy link
Collaborator Author

Thanks for this fix! BTW, I'm doing smilar things in pr #15099. Do you think adding a defensive check before importing trion, and adding dummy decorators for triton is a good way to avoid triton breaks like this? Please let me know your suggestions, thx!

This would be a much better solution. Our code is riddled with delayed imports for this reason.

@robertgshaw2-redhat
Copy link
Collaborator Author

Thanks for this fix! BTW, I'm doing smilar things in pr #15099. Do you think adding a defensive check before importing Triton, and adding dummy decorators for triton is a good way to avoid triton breaks like this? Please let me know your suggestions, thx!

Thanks @MengqingCao - will take a look

lengrongfu pushed a commit to lengrongfu/vllm that referenced this pull request Apr 2, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
kylesayrs pushed a commit to neuralmagic/vllm that referenced this pull request Apr 2, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Alex4210987 pushed a commit to LeiWang1999/vllm-bitblas that referenced this pull request Apr 5, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Signed-off-by: xinyuxiao <xinyuxiao2024@gmail.com>
lulmer pushed a commit to lulmer/vllm that referenced this pull request Apr 7, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Signed-off-by: Louis Ulmer <ulmerlouis@gmail.com>
nishith-fujitsu pushed a commit to nishith-fujitsu/vllm that referenced this pull request Apr 9, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
@yuanfz98
Copy link

Hello @robertgshaw2-redhat ,
While you moved imports into function process_weights_after_loading, the UnquantizedFusedMoEMethod will still be used in FusedMoe

if quant_config is None:

and Non-CUDA backends will run into triton import.

@MengqingCao
Copy link
Contributor

Hello @robertgshaw2-redhat , While you moved imports into function process_weights_after_loading, the UnquantizedFusedMoEMethod will still be used in FusedMoe

if quant_config is None:

and Non-CUDA backends will run into triton import.

Could you try with #15099, which is a pr addressing the triton import issue on NON-CUDA devices. Feel free to review this pr or feed back your result, thanks!

lk-chen pushed a commit to lk-chen/vllm that referenced this pull request Apr 29, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
shreyankg pushed a commit to shreyankg/vllm that referenced this pull request May 3, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
RichardoMrMu pushed a commit to RichardoMrMu/vllm that referenced this pull request May 12, 2025
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working ready ONLY add when PR is ready to merge/full CI is needed tpu Related to Google TPUs
Projects
None yet
Development

Successfully merging this pull request may close these issues.

7 participants