[Feature] Support Hierarchical Group-Cast Collective Communication #26

Strivin0311 · 2025-05-22T17:49:00Z

Supported 1D/2D cp device mesh as an interface argument.
Supported hierarchical group-cast collective:
- It is based upon all2all-v and 2d cp device mesh (inter-node x intra-node)
- We add an environment variable MAGI_ATTENTION_HIERARCHICAL_COMM to control whether to enable hierarchical group-cast.
- To enable it, you only need to set export MAGI_ATTENTION_HIERARCHICAL_COMM=1 and pass a 2D cp device mesh.
- You can use if magi_attention.is_hierarchical_comm_enable(): to check if it is enabled.
- NOTE: it had better been enabled when magi_attention.is_cuda_device_max_connections_one() == False.
- The high_bandwidth_domain_size > 1 is also forbidden when it's enabled.
Optimized and polished the code in comm/primitive/utils.py to reduce runtime overhead for range ops.
Enabled range_fill_ kernel for out_zero_fill operation in each attn_fwd_partial call.

…imitive/utils

…om interface, and added it to unitest

… pipeline

…-intra step in hier comm

…s, with all the test cases updated along

…in test_pipeline(_sdpa)

Strivin0311 force-pushed the feature/support_hier_comm branch from afb0470 to a52b290 Compare May 23, 2025 11:19

Strivin0311 added 16 commits June 3, 2025 05:40

updated the docstring for group-collective and fixed some typos in pr…

281d2db

…imitive/utils

updated some utils funcs

f7e3944

impled hier comm for group-cast, added intra/inter cp group passed fr…

0c47a7e

…om interface, and added it to unitest

added nvtx for hier comm three all2allv

1220d05

added side stream but not used; fixed a profile test case bug in test…

791d0c6

… pipeline

optimized the code in primitive/utils to reduce runtime overhead

82809e9

optimized the code in primitive/utils to reduce cpu overhead

801c0e9

added side stream for post-intra step in hier group-cast

0472189

assigned each comm a side stream for post-intra step in hier comm

f879ee9

added range gather output buffer arg; pre-init output buffer for post…

5863ea3

…-intra step in hier comm

added cp device mesh and deleted intra/inter group from all interface…

3a781dc

…s, with all the test cases updated along

used the range_fill_ for out_zero_fill

41406e6

passed all the unitest and deleted a debug print

e8ff7db

reorganized the comp-comm order and merge all error message together …

dd0af5d

…in test_pipeline(_sdpa)

added a forgotten try-catch in test pipeline

e147663

deleted a debug code

4f9313e

Strivin0311 force-pushed the feature/support_hier_comm branch from 1ad3c8e to 4f9313e Compare June 3, 2025 05:42

littsk merged commit f5089c7 into dev Jun 4, 2025
6 checks passed

littsk deleted the feature/support_hier_comm branch June 4, 2025 02:48

Provide feedback