adapt tensorrt llm custom all reduce to sgl-kernel (#2481)

Co-authored-by: Yineng Zhang <me@zhyncs.com>
2024-12-15 13:15:59 +08:00
parent 5f2595be43
commit e04d3f2897
13 changed files with 872 additions and 32 deletions
--- a/sgl-kernel/src/sgl-kernel/init.py
+++ b/sgl-kernel/src/sgl-kernel/init.py
@@ -1,3 +1,8 @@
-from .ops import warp_reduce
+from .ops import custom_dispose, custom_reduce, init_custom_reduce, warp_reduce

-__all__ = ["warp_reduce"]
+__all__ = [
+    "warp_reduce",
+    "init_custom_reduce",
+    "custom_dispose",
+    "custom_reduce",
+]