refactor qwen moe code, use communicator to support tp+dp (#6581)

This commit is contained in:
Yi Zhang
2025-05-25 23:01:10 -07:00
committed by GitHub
parent fc419b62e8
commit 65f091310c
5 changed files with 79 additions and 380 deletions
+8
View File
@@ -2026,6 +2026,14 @@ class DeepEPMode(Enum):
return DeepEPMode.normal
def is_non_idle_and_non_empty(forward_mode, hidden_states):
return (
(forward_mode is not None)
and not forward_mode.is_idle()
and hidden_states.shape[0] > 0
)
def fast_topk(values, topk, dim):
if topk == 1:
# Use max along the specified dimension to get both value and index