EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
yiakwy-xpu-ml-framework-team
2025-03-17 09:16:55 +08:00
committed by GitHub
parent f5bbf6037d
commit 9b8333d992
2 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
sgl-kernel/csrc/torch_extension_rocm.cc
View File

@@ -61,6 +61,10 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! " "moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! token_cnts_buffer, Tensor! cumsum_buffer) -> ()"); "experts_ids, Tensor! num_tokens_post_pad, Tensor! token_cnts_buffer, Tensor! cumsum_buffer) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size); m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output) -> ()");
m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
} }
REGISTER_EXTENSION(common_ops) REGISTER_EXTENSION(common_ops)

1
sgl-kernel/setup_rocm.py
View File

@@ -41,6 +41,7 @@ include_dirs = [
sources = [ sources = [
"csrc/allreduce/custom_all_reduce.hip", "csrc/allreduce/custom_all_reduce.hip",
"csrc/moe/moe_align_kernel.cu", "csrc/moe/moe_align_kernel.cu",
"csrc/moe/moe_topk_softmax_kernels.cu",
"csrc/torch_extension_rocm.cc", "csrc/torch_extension_rocm.cc",
] ]