[FusedMoE] Remove cuda hard-code in dual stream execution #27397
Conversation
Signed-off-by: Icey <1790571317@qq.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request successfully removes hard-coded CUDA dependencies from the FusedMoE layer by replacing CUDA-specific stream operations with platform-agnostic abstractions. This is a valuable change for improving hardware portability. However, I've identified a critical issue that could lead to a runtime crash on TPU platforms due to the new implementation. I have provided a suggested fix to address this potential crash.
|
I hope this PR can be included in the next version. |
| # and other execution modes | ||
| if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM: | ||
| logger.info_once("Disabling MoE shared_experts cuda stream") | ||
| if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM or current_platform.is_tpu(): |
There was a problem hiding this comment.
why add tpu check here?
There was a problem hiding this comment.
tpu using xla and the namespace is different from torch.cuda.xxx, maybe we could use torch_xla.core.xla_model for __getattr__ of tpu backend. not sure of that, maybe @DarkLight1337 could help to take a look?
4 participants
Purpose
Remove cuda hard-code in dual stream execution of FusedMoE
Test Plan
Test through the exsiting tests
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.