From 84ab32a2a0823a7f75c439f36b96e3505ec03d08 Mon Sep 17 00:00:00 2001 From: Yi Zhong <207368749+vincentzed@users.noreply.github.com> Date: Fri, 9 Jan 2026 16:53:06 -0500 Subject: [PATCH] Reduce some small cpu overhead in stream fetch (#16587) Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com> --- python/sglang/srt/models/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index e5211e833..ea2983edc 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -26,7 +26,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.utils import is_cuda +from sglang.srt.utils import get_current_device_stream_fast, is_cuda from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: @@ -246,7 +246,7 @@ def apply_qk_norm( return q, k if alt_stream is not None and get_is_capture_mode(): - current_stream = torch.cuda.current_stream() + current_stream = get_current_device_stream_fast() alt_stream.wait_stream(current_stream) q_by_head = q.reshape(-1, head_dim) q_by_head = q_norm(q_by_head)