fix: only enable moe_align_block_size for now (#2590)
This commit is contained in:
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user