Enable custom AR for AMD GPUs and maintain it in sgl-kernel (#3406)

2025-03-02 15:19:06 -08:00
parent d3fe9bae56
commit 9cf4077294
9 changed files with 1282 additions and 195 deletions
--- a/sgl-kernel/setup_rocm.py
+++ b/sgl-kernel/setup_rocm.py
@@ -44,6 +44,7 @@ include_dirs = [
 sources = [
    "src/sgl-kernel/torch_extension_rocm.cc",
    "src/sgl-kernel/csrc/moe_align_kernel.cu",
+    "src/sgl-kernel/csrc/custom_all_reduce.hip",
 ]

 cxx_flags = ["-O3"]