feature: kernel barrier and fast compile for grpcoll (#223)

* add kernel barrier for groupcoll

* split compile

* refactor grpcoll for fast compile

* code polish

* lint and update contributing.md

* resolve issues

* fix ffa build

* fix ffa build

* fix ffa build

* fix test

Author: littsk

.gitignore

@@ -1,6 +1,7 @@
 # magi_attention
 magi_attention/_version.py
 magi_attention/flex_flash_attn*
+magi_attention/csrc/comm/grpcoll/instantiations/
 *.nsys-rep
 *.ncu-rep
 

CONTRIBUTING.md

@@ -65,3 +65,18 @@ pre-commit run -a
 
 > [!NOTE]
 > Code format checking will be automatically executed when you commit your changes.
+
+
+### Type Stubs (C++ Extension)
+
+If you modify the C++ extension (`magi_attn_ext`), please remember to regenerate the Python type stubs (`.pyi` files). This ensures that static type checkers (like MyPy) and IDEs can correctly recognize the updated C++ signatures.
+
+Ensure the extension is installed in your environment, then run:
+
+```bash
+pybind11-stubgen magi_attention.magi_attn_ext -o .
+```
+
+> [!IMPORTANT]
+> Failure to update stubs after modifying C++ code may cause type checking errors during CI.
+```

exps/dist_attn/run_benchmark.py

@@ -508,7 +508,7 @@ def run_magi_attn(
     num_sms = int(getattr(ATTN_CONFIG, "num_sms", 24))
     nvl_chunk_size = int(getattr(ATTN_CONFIG, "nvl_chunk_size", 8))
     nvl_buffer_size = int(getattr(ATTN_CONFIG, "nvl_buffer_size", 256))
-    rdma_chunk_size = int(getattr(ATTN_CONFIG, "rdma_chunk_size", 4))
+    rdma_chunk_size = int(getattr(ATTN_CONFIG, "rdma_chunk_size", 16))
     rdma_buffer_size = int(getattr(ATTN_CONFIG, "rdma_buffer_size", 128))
     num_nvl_bytes = int(getattr(ATTN_CONFIG, "num_nvl_bytes", int(3e9)))  # ~3GB
     # only valid for internode

exps/grpcoll/test_internode_grpcoll.py

@@ -47,15 +47,15 @@
 from magi_attention.comm.primitive.grpcoll._buffer import GrpCollBuffer
 from magi_attention.comm.primitive.grpcoll._config import GrpCollConfig
 from magi_attention.comm.primitive.grpcoll._handle import GrpCollInterHandle
-from magi_attention.comm.primitive.grpcoll._mgr import grpcoll_mgr
+from magi_attention.comm.primitive.grpcoll._mgr import grpcoll_buffer_mgr
 from magi_attention.comm.primitive.grpcoll.utils import (
     get_a2av_perm_idxs_from_group_cast_meta,
     get_native_group_cast_meta,
     get_num_rdma_recv_tokens,
     transfer_splits_and_dst_idxs_to_t2r_idx,
     unpermute_output,
 )
-from magi_attention.common.enum import GroupReduceOp
+from magi_attention.common.enum import GroupReduceOp, GrpCollBufferName
 from magi_attention.testing.precision import assert_close
 from magi_attention.utils import pad_and_pack_tensors, setup_dist_env
 
@@ -1547,7 +1547,7 @@ def test_loop(args: argparse.Namespace):
     assert num_local_ranks == 8 and num_ranks > 8
 
     # set grpcoll config
-    use_grpcoll_mgr = True
+    use_grpcoll_buffer_mgr = True
     if args.test_ll_compatibility:
         ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
         if local_rank == 0:
@@ -1591,13 +1591,13 @@ def test_loop(args: argparse.Namespace):
         explicitly_destroy=True,
     )
 
-    if use_grpcoll_mgr:
-        grpcoll_mgr.register_buffer(
+    if use_grpcoll_buffer_mgr:
+        grpcoll_buffer_mgr.initialize(
             group=group,
             config=buffer_config,
             **extra_buffer_kwargs,
         )
-        buffer = grpcoll_mgr.get_buffer(group)
+        buffer = grpcoll_buffer_mgr.get_buffer(GrpCollBufferName.GroupCastDefault)
     else:
         buffer_args = buffer_config.to_buffer_args()
         buffer_args.update(extra_buffer_kwargs)
@@ -1619,8 +1619,8 @@ def test_loop(args: argparse.Namespace):
     )
 
     # Destroy the buffer runtime
-    if use_grpcoll_mgr:
-        grpcoll_mgr.release_buffer(group)
+    if use_grpcoll_buffer_mgr:
+        grpcoll_buffer_mgr.release_buffer(GrpCollBufferName.GroupCastDefault)
     else:
         buffer.destroy()
         dist.barrier()

exps/grpcoll/test_intranode_grpcoll.py

@@ -47,14 +47,14 @@
 from magi_attention.comm.primitive.grpcoll._buffer import GrpCollBuffer
 from magi_attention.comm.primitive.grpcoll._config import GrpCollConfig
 from magi_attention.comm.primitive.grpcoll._handle import GrpCollIntraHandle
-from magi_attention.comm.primitive.grpcoll._mgr import grpcoll_mgr
+from magi_attention.comm.primitive.grpcoll._mgr import grpcoll_buffer_mgr
 from magi_attention.comm.primitive.grpcoll.utils import (
     get_a2av_perm_idxs_from_group_cast_meta,
     get_native_group_cast_meta,
     transfer_splits_and_dst_idxs_to_t2r_idx,
     unpermute_output,
 )
-from magi_attention.common.enum import GroupReduceOp
+from magi_attention.common.enum import GroupReduceOp, GrpCollBufferName
 from magi_attention.utils import pad_and_pack_tensors
 
 # isort: split
@@ -1424,7 +1424,7 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
     rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
 
     # set grpcoll config
-    use_grpcoll_mgr = True
+    use_grpcoll_buffer_mgr = True
     test_ll_compatibility, num_rdma_bytes = False, 0
     if test_ll_compatibility:
         ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
@@ -1466,13 +1466,13 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
         explicitly_destroy=True,
     )
 
-    if use_grpcoll_mgr:
-        grpcoll_mgr.register_buffer(
+    if use_grpcoll_buffer_mgr:
+        grpcoll_buffer_mgr.initialize(
             group=group,
             config=buffer_config,
             **extra_buffer_kwargs,
         )
-        buffer = grpcoll_mgr.get_buffer(group)
+        buffer = grpcoll_buffer_mgr.get_buffer(GrpCollBufferName.GroupCastDefault)
     else:
         buffer_args = buffer_config.to_buffer_args()
         buffer_args.update(extra_buffer_kwargs)
@@ -1492,8 +1492,8 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
     )
 
     # Destroy the buffer runtime
-    if use_grpcoll_mgr:
-        grpcoll_mgr.release_buffer(group)
+    if use_grpcoll_buffer_mgr:
+        grpcoll_buffer_mgr.release_buffer(GrpCollBufferName.GroupCastDefault)
     else:
         buffer.destroy()
         dist.barrier()

magi_attention/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import importlib.util
+import logging
 import os
 import warnings
 
@@ -22,6 +23,36 @@
     init_dist_attn_runtime_mgr,
 )
 
+try:
+    from . import magi_attn_ext  # type: ignore[attr-defined]  # noqa: F401
+except ImportError as e:
+    warnings.warn(
+        f"Failed to import magi_attn_ext extension module. "
+        f"Please make sure MagiAttention is properly installed. "
+        f"Original error message: {e}"
+    )
+
+try:
+    from . import magi_attn_comm  # type: ignore[attr-defined]  # noqa: F401
+except ImportError as e:
+    warnings.warn(
+        f"Failed to import magi_attn_comm extension module. "
+        f"Please make sure MagiAttention is properly installed. "
+        f"Original error message: {e}"
+    )
+
+try:
+    from . import (  # type: ignore[attr-defined] # noqa: F401
+        flexible_flash_attention_utils_cuda,
+    )
+except ImportError as e:
+    warnings.warn(
+        f"Failed to import flexible_flash_attention_utils_cuda extension module. "
+        f"Please make sure MagiAttention is properly installed. "
+        f"Original error message: {e}"
+    )
+
+
 if importlib.util.find_spec("magi_attention._version") is None:
     warnings.warn(
         "You are using magi_attention without installing it. This may cause some unexpected errors."
@@ -34,6 +65,13 @@
 
 __version__: str | None = version
 
+# Initialize a logger specific to this module/namespace
+logger = logging.getLogger(__name__)
+
+# Add a NullHandler to prevent logging warnings ("No handlers could be found...")
+# if the application using this library hasn't configured logging.
+logger.addHandler(logging.NullHandler())
+
 
 def is_sanity_check_enable() -> bool:
     """
@@ -123,4 +161,7 @@ def dist_attn_runtime_dict_size() -> int:
     "config",
     "comm",
     "functional",
+    "magi_attn_ext",
+    "magi_attn_comm",
+    "flexible_flash_attention_utils_cuda",
 ]

magi_attention/comm/primitive/grpcoll/_buffer.py

@@ -383,6 +383,7 @@ def group_cast(
         post_perm_idx: torch.Tensor | None = None,
         config: GrpCollConfig | None = None,
         previous_event: EventOverlap | None = None,
+        kernel_barrier=None,
         async_op: bool = False,
         allocate_on_comm_stream: bool = False,
         cast_lse: bool = False,
@@ -495,6 +496,7 @@ def group_cast(
                 is_token_in_rank=is_token_in_rank,
                 post_perm_idx=post_perm_idx,
                 previous_event=previous_event,
+                kernel_barrier=kernel_barrier,
                 async_op=async_op,
                 allocate_on_comm_stream=allocate_on_comm_stream,
                 cast_lse=cast_lse,
@@ -514,6 +516,7 @@ def group_cast(
             is_token_in_rank=is_token_in_rank,
             post_perm_idx=post_perm_idx,
             previous_event=previous_event,
+            kernel_barrier=kernel_barrier,
             async_op=async_op,
             allocate_on_comm_stream=allocate_on_comm_stream,
             cast_lse=cast_lse,
@@ -531,6 +534,7 @@ def group_reduce(
         pre_perm_idx: torch.Tensor | None = None,
         config: GrpCollConfig | None = None,
         previous_event: EventOverlap | None = None,
+        kernel_barrier=None,
         async_op: bool = False,
         allocate_on_comm_stream: bool = False,
         comm_dtype: torch.dtype | None = None,
@@ -625,6 +629,7 @@ def group_reduce(
                 acc_reduce=acc_reduce,
                 pre_perm_idx=pre_perm_idx,
                 previous_event=previous_event,
+                kernel_barrier=kernel_barrier,
                 async_op=async_op,
                 allocate_on_comm_stream=allocate_on_comm_stream,
                 comm_dtype=comm_dtype,
@@ -643,6 +648,7 @@ def group_reduce(
             acc_reduce=acc_reduce,
             pre_perm_idx=pre_perm_idx,
             previous_event=previous_event,
+            kernel_barrier=kernel_barrier,
             async_op=async_op,
             allocate_on_comm_stream=allocate_on_comm_stream,
             comm_dtype=comm_dtype,
@@ -661,6 +667,7 @@ def _intranode_group_cast(
         is_token_in_rank: torch.Tensor | None = None,
         post_perm_idx: torch.Tensor | None = None,
         previous_event: EventOverlap | None = None,
+        kernel_barrier=None,
         async_op: bool = False,
         allocate_on_comm_stream: bool = False,
         cast_lse: bool = False,
@@ -747,6 +754,7 @@ def _intranode_group_cast(
             post_perm_idx,
             config.to_kernel_config(),
             getattr(previous_event, "event", None),
+            kernel_barrier,
             async_op,
             allocate_on_comm_stream,
         )
@@ -791,6 +799,7 @@ def _intranode_group_reduce(
         acc_reduce: bool = False,
         pre_perm_idx: torch.Tensor | None = None,
         previous_event: EventOverlap | None = None,
+        kernel_barrier=None,
         async_op: bool = False,
         allocate_on_comm_stream: bool = False,
         comm_dtype: torch.dtype | None = None,
@@ -843,6 +852,7 @@ def _intranode_group_reduce(
             pre_perm_idx,
             config.to_kernel_config(),
             getattr(previous_event, "event", None),
+            kernel_barrier,
             async_op,
             allocate_on_comm_stream,
             reduce_op,
@@ -873,6 +883,7 @@ def _internode_group_cast(
         is_token_in_rank: torch.Tensor | None = None,
         post_perm_idx: torch.Tensor | None = None,
         previous_event: EventOverlap | None = None,
+        kernel_barrier=None,
         async_op: bool = False,
         allocate_on_comm_stream: bool = False,
         cast_lse: bool = False,
@@ -975,6 +986,7 @@ def _internode_group_cast(
             post_perm_idx,
             config.to_kernel_config(),
             getattr(previous_event, "event", None),
+            kernel_barrier,
             async_op,
             allocate_on_comm_stream,
         )
@@ -1023,6 +1035,7 @@ def _internode_group_reduce(
         acc_reduce: bool = False,
         pre_perm_idx: torch.Tensor | None = None,
         previous_event: EventOverlap | None = None,
+        kernel_barrier=None,
         async_op: bool = False,
         allocate_on_comm_stream: bool = False,
         comm_dtype: torch.dtype | None = None,
@@ -1078,6 +1091,7 @@ def _internode_group_reduce(
             pre_perm_idx,
             config.to_kernel_config(),
             getattr(previous_event, "event", None),
+            kernel_barrier,
             async_op,
             allocate_on_comm_stream,
             reduce_op,

magi_attention/comm/primitive/grpcoll/_config.py

@@ -42,7 +42,7 @@ class GrpCollConfig:
     num_sms: int = 24
     nvl_chunk_size: int = 8
     nvl_buffer_size: int = 256
-    rdma_chunk_size: int = 4
+    rdma_chunk_size: int = 16
     rdma_buffer_size: int = 128
 
     # for buffer initialization