deepinfer-eplb-trigger

deepinfer/deepinfer/model_executor/eplb/__init__.py · a679bb8 · L146–165 · 固定 commit 上游

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

146    def _update_rearrangement(self, is_prompt: bool) -> None:
147        if self._rearrangement_interval == 0:
148            return
149        if is_prompt and self._skip_rearrangement_on_prefill:
150            return
151
152        window_index = self._expert_count_window_index
153        self._expert_count_window[window_index].copy_(self.expert_count)
154        self._expert_count_window_index = (window_index + 1) % self._expert_count_window.size(0)
155
156        if is_prompt and self._schedule_rearrangement_after_prefill:
157            # 下一个非 prefill step 立即 rearrange。
158            self._rearrangement_step = self._rearrangement_interval - 1
159            return
160
161        self._rearrangement_step += 1
162        if self._rearrangement_step < self._rearrangement_interval:
163            return
164        self._rearrangement_step = 0
165        self._rearrange_experts()