From 2dccecf43261207aaf5a8da7a92f5d1ae3f52e5b Mon Sep 17 00:00:00 2001
From: Yineng Zhang <me@zhyncs.com>
Date: Thu, 26 Dec 2024 16:56:59 +0800
Subject: [PATCH] fix: only enable moe_align_block_size for now (#2590)

---
 sgl-kernel/pyproject.toml                 |  2 +-
 sgl-kernel/src/sgl-kernel/__init__.py     | 12 +-----------
 sgl-kernel/src/sgl-kernel/ops/__init__.py | 20 --------------------
 3 files changed, 2 insertions(+), 32 deletions(-)

diff --git a/sgl-kernel/pyproject.toml b/sgl-kernel/pyproject.toml
index ebac33e9a..0ddf3b0ef 100644
--- 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
index b0aa791b1..1019896fe 100644
--- 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
index 1ca551b6b..a620f58a5 100644
--- 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(