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

[Feature] Add a fast-topk to sgl-kernel for DeepSeek v3.2 (#11194)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
DarkSharpness
2025-10-06 01:19:03 +08:00
committed by GitHub
parent 4cb5a5235e
commit e0b2d3eebe
7 changed files with 588 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
sgl-kernel/csrc/common_extension.cc
View File

@@ -107,6 +107,13 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("concat_mla_absorb_q(Tensor a, Tensor b, Tensor! out) -> ()");
m.impl("concat_mla_absorb_q", torch::kCUDA, &concat_mla_absorb_q);
m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths) -> ()");
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) -> ()");
m.impl("fast_topk_transform_fused", torch::kCUDA, &fast_topk_transform_interface);
/*
* From csrc/gemm
*/