Skip to content

[Feature]: Remove torch.distributed from template_moe all_to_all #11568

@greg-kwasniewski1

Description

@greg-kwasniewski1

🚀 The feature, motivation and pitch

Though never used, the path from torch_moe -> template_moe -> _template_moe_alltoall fallbacks to naive all-to-all using torch.distributed. Replace it with tensorrt_llm/_torch/auto_deploy/distributed/common.py

Alternatives

No response

Additional context

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Metadata

Metadata

Labels

feature requestNew feature or request. This includes new model, dtype, functionality support

Type

No type

Projects

Status

Backlog

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions