Add intel_amx backend for Radix Attention for CPU (#6408)
Co-authored-by: Chunyuan WU <chunyuan.wu@intel.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
This commit is contained in:
co-authored by
Chunyuan WU
Thien Tran
parent
e39bca0756
commit
888cb175a6
@@ -2225,3 +2225,21 @@ def bind_or_assign(target, source):
|
||||
return target
|
||||
else:
|
||||
return source
|
||||
|
||||
|
||||
def support_triton(backend: str) -> bool:
|
||||
return backend not in ["torch_native", "intel_amx"]
|
||||
|
||||
|
||||
try:
|
||||
import sgl_kernel
|
||||
|
||||
is_intel_amx_backend_available = hasattr(
|
||||
torch.ops.sgl_kernel, "convert_weight_packed"
|
||||
)
|
||||
except:
|
||||
is_intel_amx_backend_available = False
|
||||
|
||||
|
||||
def cpu_has_amx_support():
|
||||
return torch._C._cpu._is_amx_tile_supported() and is_intel_amx_backend_available
|
||||
|
||||
Reference in New Issue
Block a user