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

[AMD] Support fast_topk kernels in sgl-kernel (#15172)

Browse Source
This commit is contained in:
Hubert Lu
2025-12-19 22:19:09 -08:00
committed by GitHub
parent 6468cb5823
commit 51e2eaa458
4 changed files with 47 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension_rocm.cc
+14 -1
View File
@@ -20,7 +20,7 @@ limitations under the License.
TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
/*
* From csrc/activation
* From csrc/elementwise
*/
m.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
@@ -34,6 +34,19 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
m.def("gelu_quick(Tensor! out, Tensor input) -> ()");
m.impl("gelu_quick", torch::kCUDA, &gelu_quick);
m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths, Tensor? row_starts) -> ()");
m.impl("fast_topk", torch::kCUDA, &fast_topk_interface);
m.def(
"fast_topk_transform_fused(Tensor score, Tensor lengths, Tensor dst_page_table, Tensor src_page_table, Tensor "
"cu_seqlens_q, Tensor? row_starts) -> ()");
m.impl("fast_topk_transform_fused", torch::kCUDA, &fast_topk_transform_interface);
m.def(
"fast_topk_transform_ragged_fused(Tensor score, Tensor lengths, Tensor topk_indices_ragged, Tensor "
"topk_indices_offset, Tensor ? row_starts) -> ()");
m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface);
/*
* From csrc/allreduce
*/