[MUSA][4/N] Add common device utilities, distributed backend, and custom op wiring (#17246)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-01-28 23:13:24 -08:00
committed by GitHub
parent 9409c43593
commit d3cdee0a04
7 changed files with 109 additions and 11 deletions
@@ -20,6 +20,7 @@ from sglang.srt.utils import (
is_cpu,
is_cuda,
is_hip,
is_musa,
is_npu,
is_xpu,
)
@@ -31,6 +32,7 @@ _is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_is_xpu = is_xpu()
_is_musa = is_musa()
if _is_cuda:
from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace
@@ -119,6 +121,7 @@ class RotaryEmbedding(MultiPlatformOp):
and not (_is_cpu)
and not (_is_xpu)
and not (_is_npu)
and not (_is_musa)
):
if _is_cuda or _is_hip:
from sgl_kernel import rotary_embedding