Fix sampling for speculative decoding & simplify kernels (#7207)

2025-06-16 03:28:30 -07:00
parent b1286a116a
commit cfceb83d05
11 changed files with 124 additions and 79 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -72,6 +72,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.version import __version__

 build_tree_kernel = (