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

use fast stream instead of torch.cuda.current_stream in llama 4 shared experts overlap (#12811)

Browse Source
This commit is contained in:
b8zhong
2025-11-08 15:04:37 -08:00
committed by GitHub
parent 44f594d832
commit 49653c8896
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/models/llama4.py
View File

@@ -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