Unverified Commit f4674df6 authored by ZhichenJiang's avatar ZhichenJiang Committed by GitHub
Browse files

support idle batch for TBO (#8233)

parent d40846d4
......@@ -341,15 +341,18 @@ class TboDPAttentionPreparer:
@staticmethod
def _compute_global_forward_mode(forward_modes):
converted_forward_modes = [
ForwardMode.DECODE.value if x == ForwardMode.IDLE.value else x
for x in forward_modes
forward_modes_excluding_idle = [
x for x in forward_modes if x != ForwardMode.IDLE.value
]
if not forward_modes_excluding_idle:
return ForwardMode.IDLE, False
forward_mode_agree = TboDPAttentionPreparer._is_all_same(
converted_forward_modes
forward_modes_excluding_idle
)
global_forward_mode = (
ForwardMode(converted_forward_modes[0]) if forward_mode_agree else None
ForwardMode(forward_modes_excluding_idle[0]) if forward_mode_agree else None
)
return global_forward_mode, forward_mode_agree
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment