fix: only enable moe_align_block_size for now (#2590)

sgl-project · Dec 26, 2024 · 2dccecf · 2dccecf
1 parent 75ad0a1
commit 2dccecf
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 32 deletions.
diff --git a/sgl-kernel/pyproject.toml b/sgl-kernel/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sgl-kernel"
-version = "0.0.2.post8"
+version = "0.0.2.post9"
 description = "Kernel Library for SGLang"
 readme = "README.md"
 requires-python = ">=3.8"

diff --git a/sgl-kernel/src/sgl-kernel/__init__.py b/sgl-kernel/src/sgl-kernel/__init__.py
@@ -1,15 +1,5 @@
-from .ops import (
-    custom_dispose,
-    custom_reduce,
-    init_custom_reduce,
-    moe_align_block_size,
-    warp_reduce,
-)
+from .ops import moe_align_block_size
 
 __all__ = [
-    "warp_reduce",
-    "init_custom_reduce",
-    "custom_dispose",
-    "custom_reduce",
     "moe_align_block_size",
 ]
diff --git a/sgl-kernel/src/sgl-kernel/ops/__init__.py b/sgl-kernel/src/sgl-kernel/ops/__init__.py
@@ -1,24 +1,4 @@
-from .custom_reduce_cuda import all_reduce as _all_reduce
-from .custom_reduce_cuda import dispose as _dispose
-from .custom_reduce_cuda import init_custom_ar as _init_custom_ar
 from .moe_align_block_size import moe_align_block_size as _moe_align_block_size
-from .warp_reduce_cuda import reduce as _reduce
-
-
-def warp_reduce(input_tensor):
-    return _reduce(input_tensor)
-
-
-def init_custom_reduce(rank_id, num_devices, buffers, barrier_in, barrier_out):
-    return _init_custom_ar(rank_id, num_devices, buffers, barrier_in, barrier_out)
-
-
-def custom_dispose(fa):
-    _dispose(fa)
-
-
-def custom_reduce(fa, inp, out):
-    _all_reduce(fa, inp, out)
 
 
 def moe_align_block_size(