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/gemm.py
+++ b/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(),
    )