From 49653c88964f7818fc8c02667628252e0f883b4e Mon Sep 17 00:00:00 2001 From: b8zhong Date: Sat, 8 Nov 2025 15:04:37 -0800 Subject: [PATCH] use fast stream instead of torch.cuda.current_stream in llama 4 shared experts overlap (#12811) --- python/sglang/srt/models/llama4.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index bbcdebfb9..ca46534a7 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -58,6 +58,7 @@ from sglang.srt.utils import ( is_cuda, make_layers, ) +from sglang.srt.utils.common import get_current_device_stream_fast _is_cuda = is_cuda() @@ -164,7 +165,7 @@ class Llama4MoE(nn.Module): def _forward_core_shared_routed_overlap(self, hidden_states): alt_stream = _get_or_create_alt_stream(self.device_module) - alt_stream.wait_stream(self.device_module.current_stream()) + alt_stream.wait_stream(get_current_device_stream_fast()) shared_out = self.shared_expert(hidden_states) @@ -173,7 +174,7 @@ class Llama4MoE(nn.Module): router_logits, _ = self.router(hidden_states) topk_output = self.topk(hidden_states, router_logits) routed_out = self.experts(hidden_states, topk_output) - self.device_module.current_stream().wait_stream(alt_stream) + get_current_device_stream_fast().wait_stream(alt_stream) return shared_out, routed_out