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

[sgl-kernel][CPU] add kernel for shm_allgather_into_tensor and shm_reduce_scatter_tensor (#13397)

Browse Source
This commit is contained in:
Chunyuan WU
2026-07-23 09:19:38 +08:00
committed by GitHub
parent 9a7ac3ecef
commit 60dea26077
5 changed files with 232 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/interface.cpp
+23 -1
View File
@@ -66,5 +66,27 @@ torch::Tensor shm_allgather(torch::Tensor& data, int64_t dim) {
std::vector<int64_t> result_shape = data.sizes().vec();
result_shape[dim] *= world_size;
torch::Tensor result_tensor = torch::empty(result_shape, data.options());
return all_gather(result_tensor, data, dim, numel, data_size);
return all_gather<STATE_GROUP_ALL_GATHER>(result_tensor, data, dim, numel, data_size);
}
void shm_allgather_into_tensor(torch::Tensor& output_tensor, torch::Tensor& data) {
RECORD_FUNCTION("sgl-kernel::shm_allgather_into_tensor", std::vector<c10::IValue>({data}));
auto numel = data.numel();
int data_size = numel * data.element_size();
int64_t dim = 0;
all_gather<STATE_GROUP_ALL_GATHER_INTO_TENSOR>(output_tensor, data, dim, numel, data_size);
}
void shm_reduce_scatter_tensor(at::Tensor& output_tensor, at::Tensor& data, int64_t op) {
RECORD_FUNCTION("sgl-kernel::shm_reduce_scatter_tensor", std::vector<c10::IValue>({data}));
TORCH_CHECK(op == c10d::ReduceOp::SUM, "Only torch.distributed.ReduceOp.SUM is supported");
auto numel = data.numel();
int data_size = numel * data.element_size();
reduce_scatter_outer_loop(output_tensor, data, numel, data_size);
return;
}