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/torch_extension_cpu.cpp
+10
View File
@@ -440,6 +440,12 @@ void shm_allreduce(at::Tensor& data, int64_t op);
// shared memory all_gather
at::Tensor shm_allgather(at::Tensor& data, int64_t dim);
// shared memory all_gather_into_tensor
void shm_allgather_into_tensor(at::Tensor& output_tensor, at::Tensor& data);
// shared memory reduce_scatter_tensor
void shm_reduce_scatter_tensor(at::Tensor& output_tensor, at::Tensor& data, int64_t op);
// rope
std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
at::Tensor& positions,
@@ -810,6 +816,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.impl("shm_allreduce", torch::kCPU, &shm_allreduce);
m.def("shm_allgather(Tensor data, int dim) -> Tensor");
m.impl("shm_allgather", torch::kCPU, &shm_allgather);
m.def("shm_allgather_into_tensor(Tensor(a!) output_tensor, Tensor data) -> ()");
m.impl("shm_allgather_into_tensor", torch::kCPU, &shm_allgather_into_tensor);
m.def("shm_reduce_scatter_tensor(Tensor(a!) output_tensor, Tensor data, int reduce_op) -> ()");
m.impl("shm_reduce_scatter_tensor", torch::kCPU, &shm_reduce_scatter_tensor);
// rope
m.def(