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

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
2025-10-06 01:19:03 +08:00
parent 4cb5a5235e
commit e0b2d3eebe
7 changed files with 588 additions and 1 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -309,7 +309,7 @@ from sgl_kernel.speculative import (
    tree_speculative_sampling_target_only,
    verify_tree_greedy,
 )
-from sgl_kernel.top_k import fast_topk
+from sgl_kernel.top_k import fast_topk, fast_topk_transform_fused, fast_topk_v2
 from sgl_kernel.version import __version__

 if torch.version.hip is not None: