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:
YanbingJiang
2025-06-17 22:11:50 -07:00
committed by GitHub
co-authored by Wu, Chunyuan <chunyuan.wu@intel.com> jianan-gu sdp
parent e56685ac1b
commit 094c116f7d
23 changed files with 270 additions and 56 deletions
+5 -1
View File
@@ -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