diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 541eb94ef..248839be6 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -349,11 +349,9 @@ class BailingMoESparseMoeBlock(nn.Module): num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) - DUAL_STREAM_TOKEN_THRESHOLD = 1024 if ( self.alt_stream is not None and hidden_states.shape[0] > 0 - and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD and get_is_capture_mode() ): final_hidden_states, shared_output = self.forward_normal_dual_stream( diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 104068fe5..5d93cb735 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -125,13 +125,11 @@ class KimiMoE(nn.Module): hidden_states = hidden_states.view(-1, hidden_size) shared_output = None - DUAL_STREAM_TOKEN_THRESHOLD = 1024 if ( self.alt_stream is not None and self.num_shared_experts is not None and hidden_states.shape[0] > 0 - and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD and get_is_capture_mode() ): current_stream = torch.cuda.current_stream() diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index b89c62f44..b8bfc81ef 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -349,11 +349,9 @@ class LLaDA2MoeSparseMoeBlock(nn.Module): num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) - DUAL_STREAM_TOKEN_THRESHOLD = 1024 if ( self.alt_stream is not None and hidden_states.shape[0] > 0 - and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD and get_is_capture_mode() ): final_hidden_states, shared_output = self.forward_normal_dual_stream( diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index caeae9c23..ea33e81ef 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -304,11 +304,9 @@ class Qwen2MoeSparseMoeBlock(nn.Module): if get_moe_a2a_backend().is_deepep(): return self._forward_deepep(hidden_states, forward_batch) - DUAL_STREAM_TOKEN_THRESHOLD = 1024 if ( self.alt_stream is not None and hidden_states.shape[0] > 0 - and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD and get_is_capture_mode() ): final_hidden_states, shared_output = self.forward_normal_dual_stream(