[ROCm] enable moe topk softmax in amd (#4448)

2025-03-17 09:16:55 +08:00
parent f5bbf6037d
commit 9b8333d992
2 changed files with 5 additions and 0 deletions
--- a/sgl-kernel/setup_rocm.py
+++ b/sgl-kernel/setup_rocm.py
@@ -41,6 +41,7 @@ include_dirs = [
 sources = [
    "csrc/allreduce/custom_all_reduce.hip",
    "csrc/moe/moe_align_kernel.cu",
+    "csrc/moe/moe_topk_softmax_kernels.cu",
    "csrc/torch_extension_rocm.cc",
 ]