move all get_stream in sgl_kernel to c++ to reduce the launch overhead (#12521)

2025-11-02 13:15:05 -08:00
parent c9db79117f
commit 20315697f4
20 changed files with 61 additions and 94 deletions
--- a/sgl-kernel/python/sgl_kernel/speculative.py
+++ b/sgl-kernel/python/sgl_kernel/speculative.py
@@ -1,5 +1,4 @@
 import torch
-from sgl_kernel.utils import get_cuda_stream


 def tree_speculative_sampling_target_only(
@@ -33,7 +32,6 @@ def tree_speculative_sampling_target_only(
        threshold_single,
        threshold_acc,
        deterministic,
-        get_cuda_stream(),
    )


@@ -56,7 +54,6 @@ def verify_tree_greedy(
        retrive_next_token,
        retrive_next_sibling,
        target_predict,
-        get_cuda_stream(),
    )