[cpu/arm64] support run sglang on arm64 cpu (#14867)
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user