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

[sgl-kernel] Support moe_sum_reduce cuda kernel (#10321)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Yuan Luo
2025-09-19 14:12:09 +08:00
committed by GitHub
parent ac2a723bb3
commit 616a3e20df
7 changed files with 346 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sgl-kernel/include/sgl_kernel_ops.h
View File

@@ -293,6 +293,8 @@ void moe_align_block_size(
void topk_softmax(
torch::Tensor& topk_weights, torch::Tensor& topk_indices, torch::Tensor& gating_output, bool renormalize);
void moe_sum_reduce(at::Tensor& input, at::Tensor& output, double routed_scaling_factor);
std::vector<at::Tensor> moe_fused_gate(
at::Tensor& input,
at::Tensor& bias,