biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[torch.compile] register allreduce operations as custom ops (#8526)

Browse Source
This commit is contained in:
youkaichao
2024-09-16 22:57:57 -07:00
committed by GitHub
parent ee2bceaaa6
commit 99aa4eddaf
9 changed files with 137 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
csrc/torch_bindings.cpp
View File

@@ -411,11 +411,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
"bool full_nvlink) -> int");
custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
custom_ar.def(
"should_custom_ar(Tensor inp, int max_size, int world_size, "
"bool full_nvlink) -> bool");
custom_ar.impl("should_custom_ar", torch::kCUDA, &should_custom_ar);
custom_ar.def("all_reduce_reg(int fa, Tensor inp, Tensor! out) -> ()");
custom_ar.impl("all_reduce_reg", torch::kCUDA, &all_reduce_reg);