deepinfer-moe-forward

deepinfer/deepinfer/model_executor/models/v4_modules/mlp.py · a679bb8 · L565–616 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

565    def forward(
566        self,
567        hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
568        hidden_states_bf16: torch.Tensor,
569        input_metadata: InputMetadata,
570    ) -> torch.Tensor:
571        if self.use_mega_moe:
572            return self.mega_moe_forward(hidden_states, hidden_states_bf16, input_metadata)
573
574        topk_idx, topk_weights = self.gate_forward(hidden_states_bf16, input_metadata)
575
576        try:
577            # Dispatch
578            event_overlap = self.dispatch(
579                hidden_states,
580                topk_idx,
581                topk_weights,
582                input_metadata,
583                None,
584            )
585            moe_in_hidden, _dispatched_topk_idx, dispatched_topk_weights, handle = event_overlap.current_stream_wait()
586
587            # MoE compute
588            with cuda_graph_timer.record(
589                "routed_experts",
590                enabled=(self.layer_idx == 1 and not input_metadata.is_prompt),
591                scale=eplb.current().num_moe_layers,
592            ):
593                moe_out_hidden = self.routed_experts_forward(
594                    moe_in_hidden,
595                    dispatched_topk_weights,
596                    handle.psum_num_recv_tokens_per_expert,
597                    expected_m=None,
598                )
599
600            # Combine
601            combine_pre_evt = self.ep_buffer.capture()
602            with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
603                shared_out = self.shared_forward(hidden_states)
604            event_overlap = self.combine(
605                moe_out_hidden,
606                shared_out,
607                handle,
608                combine_pre_evt,
609            )
610            moe_combined_hidden, _combined_topk_weights = event_overlap.current_stream_wait()
611
612        except Exception:
613            logger.exception(f"[rank {PG.get_rank()}] Exception during MoE forward, topk_idx.shape = {topk_idx.shape}")
614            raise
615
616        return moe_combined_hidden