Add custom communicator for trtllm_mnnvl_ar #2056
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the ✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
| group_rank: int, | ||
| device_idx: int, | ||
| is_multi_node: bool = True, | ||
| comm: Optional[CommBackend] = None, |
There was a problem hiding this comment.
How about calling it comm_backend or communicator?
cc @nvmbreughe in case you have any preference.
There was a problem hiding this comment.
comm_backend sounds good. Or even more explicit comm_backend_for_handle_transfer Besides the name I would also list some options: MpiComm()
| group_rank: int, | ||
| device_idx: int, | ||
| is_multi_node: bool = True, | ||
| comm: Optional[CommBackend] = None, |
There was a problem hiding this comment.
comm_backend sounds good. Or even more explicit comm_backend_for_handle_transfer Besides the name I would also list some options: MpiComm()
| ) | ||
|
|
||
| @torch.inference_mode() | ||
| def row_linear_residual_norm_forward( |
There was a problem hiding this comment.
Can we make this a helper function as tests/comm/test_trtllm_mnnvl_allreduce.py also uses it?
3 participants
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes