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

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from typing import List, Optional
 
 import torch
-from sgl_kernel.utils import get_cuda_stream, is_arch_support_pdl
+from sgl_kernel.utils import is_arch_support_pdl
 
 
 # These implementations extensively draw from and build upon the FlashInfer project https://github.com/flashinfer-ai/flashinfer
@@ -263,6 +263,10 @@ class FusedSetKVBufferArg:
     cache_loc: torch.Tensor
 
 
+def _view_3d(x, head_size):
+    return x.view(x.shape[0], -1, head_size)
+
+
 def apply_rope_with_cos_sin_cache_inplace(
     positions: torch.Tensor,
     query: torch.Tensor,
@@ -317,31 +321,27 @@ def apply_rope_with_cos_sin_cache_inplace(
         assert a.v_scale is None, "v_scale is not yet supported"
         assert a.cache_loc.dtype == torch.int64, f"{a.cache_loc.dtype=}"
 
-    def _view_3d(x):
-        return x.view(x.shape[0], -1, head_size)
-
     torch.ops.sgl_kernel.apply_rope_pos_ids_cos_sin_cache.default(
-        _view_3d(query),
-        _view_3d(key),
-        _view_3d(query),
-        _view_3d(key),
+        _view_3d(query, head_size),
+        _view_3d(key, head_size),
+        _view_3d(query, head_size),
+        _view_3d(key, head_size),
         cos_sin_cache,
         positions.long(),
         (not is_neox),
         enable_pdl,
-        get_cuda_stream(),
         (
-            _view_3d(fused_set_kv_buffer_arg.value)
+            _view_3d(fused_set_kv_buffer_arg.value, head_size)
             if fused_set_kv_buffer_arg is not None
             else None
         ),
         (
-            _view_3d(fused_set_kv_buffer_arg.k_buffer)
+            _view_3d(fused_set_kv_buffer_arg.k_buffer, head_size)
             if fused_set_kv_buffer_arg is not None
             else None
         ),
         (
-            _view_3d(fused_set_kv_buffer_arg.v_buffer)
+            _view_3d(fused_set_kv_buffer_arg.v_buffer, head_size)
             if fused_set_kv_buffer_arg is not None
             else None
         ),
@@ -365,7 +365,7 @@ def downcast_fp8(
     offset: int = 0,
 ) -> None:
     torch.ops.sgl_kernel.downcast_fp8(
-        k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset, get_cuda_stream()
+        k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset
     )
 
 

sgl-kernel/python/sgl_kernel/gemm.py

@@ -2,7 +2,7 @@ from typing import Optional, Tuple
 
 import torch
 from sgl_kernel.scalar_type import ScalarType
-from sgl_kernel.utils import _get_cache_buf, get_cuda_stream
+from sgl_kernel.utils import _get_cache_buf
 
 
 def awq_dequantize(
@@ -60,7 +60,6 @@ def _bmm_fp8_internal(
         B_scale,
         workspace_buffer,
         cublas_handle,
-        get_cuda_stream(),
     )
 
 

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(),
     )
 
 

sgl-kernel/python/sgl_kernel/utils.py

@@ -18,11 +18,6 @@ from typing import Dict, Tuple
 
 import torch
 
-
-def get_cuda_stream() -> int:
-    return torch.cuda.current_stream().cuda_stream
-
-
 _cache_buf: Dict[Tuple[str, torch.device], torch.Tensor] = {}