minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Reduce some small cpu overhead in stream fetch (#16587)

Browse Source 
Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
This commit is contained in:
Yi Zhong
2026-01-09 21:53:06 +00:00
committed by GitHub
parent 7066711529
commit 84ab32a2a0
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/utils.py
+2 -2
View File
@@ -26,7 +26,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.radix_attention import RadixAttention 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.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch 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 from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -246,7 +246,7 @@ def apply_qk_norm(
return q, k return q, k
if alt_stream is not None and get_is_capture_mode(): 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) alt_stream.wait_stream(current_stream)
q_by_head = q.reshape(-1, head_dim) q_by_head = q.reshape(-1, head_dim)
q_by_head = q_norm(q_by_head) q_by_head = q_norm(q_by_head)