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

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

@@ -174,6 +174,14 @@ void copy_to_gpu_no_ce(const at::Tensor& input, at::Tensor& output);
void concat_mla_k(torch::Tensor k, torch::Tensor k_nope, torch::Tensor k_rope);
void concat_mla_absorb_q(at::Tensor a, at::Tensor b, at::Tensor out);
void fast_topk_interface(at::Tensor score, at::Tensor indices, at::Tensor lengths);
void fast_topk_transform_interface(
at::Tensor score,
at::Tensor lengths,
at::Tensor dst_page_table,
at::Tensor src_page_table,
at::Tensor cu_seqlens_q);
#ifdef USE_ROCM
void gelu_quick(at::Tensor& out, const at::Tensor& input);
#endif