[10/N] MoE Refactor: reorganize deepgemm runner in DeepEPMoE (#12054)

2025-10-24 19:16:17 -07:00
parent 729b242934
commit 649949807f
11 changed files with 397 additions and 349 deletions
@@ -1005,16 +1005,14 @@ class DeepseekV2MoE(nn.Module):
                )

    def op_experts(self, state):
-        state.hidden_states_experts_output = self.experts.run_moe_core(
+        state.combine_input = self.experts.run_moe_core(
            dispatch_output=state.dispatch_output,
        )

    def op_combine_a(self, state):
        if self.ep_size > 1:
            self.experts.dispatcher.combine_a(
-                hidden_states=state.pop("hidden_states_experts_output"),
-                topk_ids=state.dispatch_output.topk_ids,
-                topk_weights=state.dispatch_output.topk_weights,
+                combine_input=state.pop("combine_input"),
                tbo_subbatch_index=state.get("tbo_subbatch_index"),
            )
            state.pop("dispatch_output")