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()