chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU][Fix CI] Solidate torch version for sgl-kernel-cpu and fix device orientation error (#17460)

Browse Source
This commit is contained in:
Zaili Wang
2026-01-22 14:04:50 +08:00
committed by GitHub
parent e2d33531f3
commit 672eb37534
5 changed files with 27 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/sglang/srt/distributed/parallel_state.py
View File

@@ -238,8 +238,21 @@ class GroupCoordinator:
self.cpu_group = None
self.local_size = get_int_env_var("LOCAL_SIZE", 0)
if is_cuda_alike():
device_id = (
0 if envs.SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS.get() else local_rank
)
self.device = torch.device(f"cuda:{device_id}")
elif _is_npu:
self.device = torch.device(f"npu:{local_rank}")
elif _is_xpu:
self.device = torch.device(f"xpu:{local_rank}")
else:
self.device = torch.device("cpu")
self.device_module = torch.get_device_module(self.device)
for ranks in group_ranks:
active_ranks = torch.ones(len(ranks), dtype=torch.int32, device="cuda")
active_ranks = torch.ones(len(ranks), dtype=torch.int32, device=self.device)
active_ranks_cpu = torch.ones(len(ranks), dtype=torch.int32)
if "mooncake" in torch_distributed_backend:
from mooncake.ep import MooncakeBackendOptions
@@ -275,17 +288,6 @@ class GroupCoordinator:
assert self.cpu_group is not None
assert self.device_group is not None
if is_cuda_alike():
device_id = (
0 if envs.SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS.get() else local_rank
)
self.device = torch.device(f"cuda:{device_id}")
elif _is_npu:
self.device = torch.device(f"npu:{local_rank}")
else:
self.device = torch.device("cpu")
self.device_module = torch.get_device_module(self.device)
# Import communicators
self.use_pynccl = use_pynccl
self.pynccl_use_current_stream = pynccl_use_current_stream