Skip to content

A unified API for the MNNVL and single-node AllReduce kernels. #2130

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

Draft
wants to merge 8 commits into
base: main
Choose a base branch
from
Draft

A unified API for the MNNVL and single-node AllReduce kernels. #2130

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a4c9da2
Added first non-working version
nvmbreughe Nov 21, 2025
1917c76
Polished the interface
nvmbreughe Nov 21, 2025
9225e52
Removed device param
nvmbreughe Nov 21, 2025
84e75e4
Updated test with legacy vs unified API
nvmbreughe Nov 21, 2025
3bb586b
Fixed unit test
nvmbreughe Nov 24, 2025
0bcc8da
Relaxed check on trtllm_ar
nvmbreughe Nov 24, 2025
0c1391d
Made metadata mandatory in unified API, added workspace check functions
nvmbreughe Nov 24, 2025
b47ade4
Merged dtype and use_fp32_lamport params
nvmbreughe Nov 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions flashinfer/comm/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,4 +39,16 @@
from .vllm_ar import register_buffer as vllm_register_buffer
from .vllm_ar import register_graph_buffers as vllm_register_graph_buffers

# Unified AllReduce Fusion API
from .allreduce import AllReduceFusionWorkspace as AllReduceFusionWorkspace
from .allreduce import MNNVLAllReduceFusionWorkspace as MNNVLAllReduceFusionWorkspace
from .allreduce import TRTLLMAllReduceFusionWorkspace as TRTLLMAllReduceFusionWorkspace
from .allreduce import allreduce_fusion as allreduce_fusion
from .allreduce import (
create_allreduce_fusion_workspace as create_allreduce_fusion_workspace,
)
from .allreduce import (
destroy_allreduce_fusion_workspace as destroy_allreduce_fusion_workspace,
)

# from .mnnvl import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo
Loading