deepinfer-moe-communication

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

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

461    def dispatch(
462        self,
463        hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
464        topk_idx: torch.Tensor,
465        topk_weights: torch.Tensor,
466        input_metadata: InputMetadata,
467        dispatch_pre_evt: EventHandle | None,
468    ) -> EventOverlap:
469        event_overlap = self.ep_buffer.dispatch(
470            x=hidden_states,
471            num_sms=envs.DEEPINFER_USERBUFFER_COMM_SMS,  # TODO: Use all SMs when DeepEP support different #SM for dispatch and combine.
472            topk_idx=topk_idx,
473            topk_weights=topk_weights,
474            cumulative_local_expert_recv_stats=self._eplb_layer().expert_count,
475            num_experts=self.num_physical_experts,
476            num_max_tokens_per_rank=self.max_num_prefilling_tokens if input_metadata.is_prompt else self.max_num_decoding_tokens,
477            expert_alignment=deep_gemm.get_mk_alignment_for_contiguous_layout(),
478            previous_event=dispatch_pre_evt,
479            async_with_compute_stream=True,
480            allocate_on_comm_stream=dispatch_pre_evt is not None,
481            do_handle_copy=False,
482            # TODO do_sync=True makes yoco compaction a cpu bottleneck,
483            # while =False costs extra GPU memory. Mask mode saves no compute and the CPU is not the bottleneck,
484            # so it is disabled only in compaction mode.
485            do_cpu_sync=input_metadata.is_prompt and envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE != "compact",
486            do_expand=True,
487            use_tma_aligned_col_major_sf=True,
488            defer_epilogue=True,
489        )
490        return event_overlap
491
492    # noinspection PyMethodMayBeStatic
493    def combine(
494        self,
495        hidden_states: torch.Tensor,
496        shared_hidden_states: torch.Tensor | None,
497        handle: EPHandle,
498        combine_pre_evt: EventHandle | None,
499    ) -> EventOverlap:
500        # Launch high-throughput chunked combine.
501        event_overlap = self.ep_buffer.combine(
502            x=hidden_states,
503            handle=handle,
504            num_sms=envs.DEEPINFER_USERBUFFER_COMM_SMS,
505            bias=shared_hidden_states,
506            previous_event=combine_pre_evt,
507            async_with_compute_stream=True,
508            allocate_on_comm_stream=combine_pre_evt is not None,
509            defer_epilogue=True,
510        )
511        return event_overlap