diff --git a/python/sglang/srt/hardware_backend/npu/cmo.py b/python/sglang/srt/hardware_backend/npu/cmo.py index 40f3b4f16..214097de6 100644 --- a/python/sglang/srt/hardware_backend/npu/cmo.py +++ b/python/sglang/srt/hardware_backend/npu/cmo.py @@ -1,6 +1,7 @@ import torch cmo_stream = None +share_stream = None def get_cmo_stream(): @@ -52,3 +53,31 @@ def wait_cmo_stream(): if stream is not None: cur_stream = torch.npu.current_stream() cur_stream.wait_stream(stream) + + +def get_share_stream(): + global share_stream + return share_stream + + +def set_share_stream(stream): + global share_stream + share_stream = stream + + +def wait_share_stream(): + stream = get_share_stream() + if stream is not None: + cur_stream = torch.npu.current_stream() + cur_stream.wait_stream(stream) + + +def shared_expert_on_independent_stream(hidden_states, forward_func): + stream = get_share_stream() + if stream is None: + stream = torch.npu.Stream() + set_share_stream(stream) + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + shared_output = forward_func(hidden_states) + return shared_output diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 2c6fd4da7..06b272022 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -94,9 +94,18 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_npu, make_layers, use_intel_amx_backend, ) + +if is_npu(): + from sglang.srt.hardware_backend.npu.cmo import ( + shared_expert_on_independent_stream, + wait_share_stream, + ) + +from sglang.srt.environ import envs from sglang.srt.utils.hf_transformers_utils import get_rope_config logger = logging.getLogger(__name__) @@ -381,11 +390,21 @@ class Qwen2MoeSparseMoeBlock(nn.Module): return shared_output def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + enable_dual_stream = ( + is_npu() + and envs.SGLANG_NPU_USE_MULTI_STREAM.get() + and forward_batch.forward_mode.is_cuda_graph() + ) shared_output = None if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - shared_output = self._forward_shared_experts(hidden_states) + if enable_dual_stream: + shared_output = shared_expert_on_independent_stream( + hidden_states.clone(), self._forward_shared_experts + ) + else: + shared_output = self._forward_shared_experts(hidden_states) topk_output = self.topk( hidden_states, router_logits, @@ -404,6 +423,8 @@ class Qwen2MoeSparseMoeBlock(nn.Module): hidden_states=hidden_states, topk_output=topk_output, ) + if enable_dual_stream: + wait_share_stream() if shared_output is not None: final_hidden_states.add_(shared_output)