minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[sgl-kernel][1/2] Fused qk_norm_rope for Qwen3-MoE (#14036)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-11-28 12:25:15 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent bce40fa217
commit e12c78aab6
7 changed files with 696 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension.cc
+8
View File
@@ -264,9 +264,17 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("shuffle_rows(Tensor input, Tensor dst2src_map, Tensor output) -> ()");
m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
m.def("apply_shuffle_mul_sum(Tensor input, Tensor output, Tensor permutation, Tensor? factors) -> ()");
m.impl("apply_shuffle_mul_sum", torch::kCUDA, &apply_shuffle_mul_sum);
m.def(
"fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
"int num_heads_k, int num_heads_v, int head_dim, float eps, "
"Tensor q_weight, Tensor k_weight, float base, "
"bool is_neox, Tensor position_ids, float factor, float low, float high, float attention_factor) -> ()");
m.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
/*
* From csrc/moe/cutlass_moe/w4a8
*/