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

[amd] Add deterministic all-reduce kernel for AMD (ROCm) (#15340)

Browse Source 
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
This commit is contained in:
sunxxuns
2025-12-18 23:36:03 -08:00
committed by GitHub
co-authored by Thomas Wang
parent 0e869f0868
commit f2d64e6782
12 changed files with 1466 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension_rocm.cc
+11
View File
@@ -51,6 +51,17 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"()");
m.impl("all_reduce_unreg", torch::kCUDA, &all_reduce_unreg);
// Deterministic all-reduce for ROCm
extern void deterministic_all_reduce_reg(int64_t _fa, torch::Tensor & inp, torch::Tensor & out);
extern void deterministic_all_reduce_unreg(
int64_t _fa, torch::Tensor & inp, torch::Tensor & reg_buffer, torch::Tensor & out);
m.def("deterministic_all_reduce_reg(int fa, Tensor inp, Tensor! out) -> ()");
m.impl("deterministic_all_reduce_reg", torch::kCUDA, &deterministic_all_reduce_reg);
m.def("deterministic_all_reduce_unreg(int fa, Tensor inp, Tensor reg_buffer, Tensor! out) -> ()");
m.impl("deterministic_all_reduce_unreg", torch::kCUDA, &deterministic_all_reduce_unreg);
m.def("dispose", &dispose);
m.def("meta_size", &meta_size);