Use dual stream for DS MoE whenever cuda graph is used (instead of with token threshold) (#9405)
This commit is contained in:
@@ -787,12 +787,13 @@ class DeepseekV2MoE(nn.Module):
|
|||||||
gemm_output_zero_allocator: BumpAllocator = None,
|
gemm_output_zero_allocator: BumpAllocator = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
if not self._enable_a2a_moe:
|
if not self._enable_a2a_moe:
|
||||||
DUAL_STREAM_TOKEN_THRESHOLD = 1024
|
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
|
||||||
|
|
||||||
if (
|
if (
|
||||||
self.alt_stream is not None
|
self.alt_stream is not None
|
||||||
and self.num_fused_shared_experts == 0
|
and self.num_fused_shared_experts == 0
|
||||||
and hidden_states.shape[0] > 0
|
and hidden_states.shape[0] > 0
|
||||||
and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD
|
and get_is_capture_mode()
|
||||||
):
|
):
|
||||||
return self.forward_normal_dual_stream(
|
return self.forward_normal_dual_stream(
|
||||||
hidden_states,
|
hidden_states,
|
||||||
|
|||||||
Reference in New Issue
Block a user