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:
YanbingJiang
2025-05-30 21:37:42 -07:00
committed by GitHub
co-authored by Chunyuan WU Thien Tran
parent e39bca0756
commit 888cb175a6
8 changed files with 185 additions and 5 deletions
+18
View File
@@ -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