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

Clean up sgl kernel (#12413)

Browse Source 
Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>
This commit is contained in:
Lianmin Zheng
2025-10-31 01:13:34 -07:00
committed by GitHub
co-authored by Byron Hsu
parent 2e48584b62
commit c0652d907b
10 changed files with 514 additions and 379 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension_rocm.cc
+9 -9
View File
@@ -70,7 +70,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
m.impl("get_meta_buffer_ipc_handle", torch::kCPU, &get_meta_buffer_ipc_handle);
// quick allreduce
#ifdef USE_ROCM
m.def(
"qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
"cast_bf2half) -> ()");
@@ -86,7 +85,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
// Max input size in bytes
m.def("qr_max_size", &qr_max_size);
#endif
/*
* From csrc/moe
@@ -97,7 +95,9 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"pad_sorted_token_ids) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
m.def("topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize) -> ()");
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, float "
"moe_softcapping, Tensor? correction_bias) -> ()");
m.impl("topk_softmax", torch::kCUDA, &topk_softmax);
/*
@@ -116,12 +116,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"()");
m.impl("build_tree_kernel_efficient", torch::kCUDA, &build_tree_kernel_efficient);
/*
* From XGrammar
*/
m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()");
m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace);
/*
* From csrc/kvcacheio
*/
@@ -171,6 +165,12 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"transfer_kv_all_layer_direct_lf_pf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int page_size) ->() ");
m.impl("transfer_kv_all_layer_direct_lf_pf", torch::kCUDA, &transfer_kv_all_layer_direct_lf_pf);
/*
* From csrc/grammar
*/
m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()");
m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace);
}
REGISTER_EXTENSION(common_ops)