fix: only enable moe_align_block_size for now (#2590)

This commit is contained in:
Yineng Zhang
2024-12-26 16:56:59 +08:00
committed by GitHub
parent 75ad0a143f
commit 2dccecf432
3 changed files with 2 additions and 32 deletions

View File

@@ -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"

View File

@@ -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",
]

View File

@@ -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(