[Feature] Integrate quick allreduce and select the best allreduce implementation (#6619)

Signed-off-by: Haoyang Li <Haoyang.Li@amd.com>
Co-authored-by: ilmarkov <imarkov@redhat.com>

sgl-kernel/python/sgl_kernel/allreduce.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 
@@ -49,6 +49,38 @@ if torch.version.hip is not None:
     def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
         return torch.ops.sgl_kernel.get_meta_buffer_ipc_handle.default(inp)
 
+    # ROCM quick allreduce
+    def init_custom_qr(
+        rank: int, world_size: int, qr_max_size: Optional[int] = None
+    ) -> int:
+        return torch.ops.sgl_kernel.init_custom_qr.default(
+            world_size, rank, qr_max_size
+        )
+
+    def qr_get_handle(fa: int) -> torch.Tensor:
+        return torch.ops.sgl_kernel.qr_get_handle.default(fa)
+
+    def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
+        torch.ops.sgl_kernel.qr_open_handles.default(fa, handles)
+
+    def qr_all_reduce(
+        fa: int,
+        profile: int,
+        inp: torch.Tensor,
+        out: torch.Tensor,
+        cast_bf162half: bool,
+    ) -> None:
+        torch.ops.sgl_kernel.qr_all_reduce.default(
+            fa, profile, inp, out, cast_bf162half
+        )
+
+    def qr_destroy(fa: int) -> None:
+        torch.ops.sgl_kernel.qr_destroy.default(fa)
+
+    def qr_max_size() -> int:
+        return torch.ops.sgl_kernel.qr_max_size.default()
+
+    # mscclpp
     def mscclpp_generate_unique_id() -> bytes:
         raise NotImplementedError()