[cpu/arm64] support run sglang on arm64 cpu (#14867)

This commit is contained in:
Yibo Cai
2026-01-11 20:27:19 +08:00
committed by GitHub
parent f9fc50acd6
commit 2f4a6addf3
5 changed files with 17 additions and 26 deletions

View File

@@ -148,6 +148,7 @@ from sglang.srt.utils import (
get_local_ip_auto,
init_custom_process_group,
is_hip,
is_host_cpu_arm64,
is_npu,
log_info_on_rank0,
monkey_patch_p2p_access_check,
@@ -178,6 +179,7 @@ from sglang.srt.weight_sync.tensor_bucket import (
_is_hip = is_hip()
_is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu_arm64 = is_host_cpu_arm64()
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import init_npu_backend
@@ -729,7 +731,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if not self.is_draft_worker:
if self.device == "cpu":
if _is_cpu_amx_available:
if _is_cpu_amx_available or _is_cpu_arm64:
# Bind OpenMP threads to CPU cores
torch.ops.sgl_kernel.init_cpu_threads_env(self.local_omp_cpuid)
@@ -743,7 +745,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
else:
logger.warning(
"init_cpu_threads_env and shared memory based AllReduce is disabled since intel amx backend is not available"
"init_cpu_threads_env and shared memory based AllReduce is disabled, only intel amx backend and arm64 are supported"
)
# Only initialize the distributed environment on the target model worker.