deepinfer-moe-gate

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

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

379        eplb_layer = self._eplb_layer()
380
381        gate_kwargs = dict(
382            num_topk=self.num_topk,
383            use_shared_as_routed=False,
384            num_shared_experts=0,
385            routed_scaling_factor=self.routed_scaling_factor,
386            ep_rank=self.ep_rank,
387            mask=input_metadata.valid_token_mask,
388            bias=self.e_score_correction_bias,
389            image_bias=self.e_score_correction_bias_vl if images_seq_mask is not None else None,
390            image_token_mask=images_seq_mask,
391            force_random=gate_force_random,
392            to_physical_map=eplb_layer.logical_to_physical_map,
393            logical_count=eplb_layer.logical_count,
394            scoring_func=self.scoring_func,
395            out=out,
396        )
397        if self.use_mega_gate:
398            topk_idx, topk_weights = deep_gemm.bf16_mega_gate(hidden_states, self.gate.weight, **gate_kwargs)
399        else:
400            topk_idx, topk_weights = moe_topk_gate(self.gate(hidden_states), **gate_kwargs)
401        return topk_idx, topk_weights