Update python API of activation, topk, norm and rope and remove vllm dependency (#6614)
Co-authored-by: Wu, Chunyuan <chunyuan.wu@intel.com> Co-authored-by: jianan-gu <jianan.gu@intel.com> Co-authored-by: sdp <sdp@gnr799219.jf.intel.com>
This commit is contained in:
co-authored by
Wu, Chunyuan <chunyuan.wu@intel.com>
jianan-gu
sdp
parent
e56685ac1b
commit
094c116f7d
@@ -1,9 +1,11 @@
|
||||
from torch import nn
|
||||
|
||||
from sglang.srt.utils import is_cuda, is_hip
|
||||
from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_hip = is_hip()
|
||||
_is_cpu = is_cpu()
|
||||
_is_cpu_amx_available = cpu_has_amx_support()
|
||||
|
||||
|
||||
class CustomOp(nn.Module):
|
||||
@@ -75,5 +77,7 @@ class CustomOp(nn.Module):
|
||||
return self.forward_cuda
|
||||
elif _is_hip:
|
||||
return self.forward_hip
|
||||
elif _is_cpu and _is_cpu_amx_available:
|
||||
return self.forward_cpu
|
||||
else:
|
||||
return self.forward_native
|
||||
|
||||
Reference in New Issue
Block a user