[TBO] fix cuda graph intermittently becomes disabled bug (#18320)
This commit is contained in:
@@ -440,18 +440,23 @@ class TboDPAttentionPreparer:
|
|||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _compute_global_forward_mode(forward_modes):
|
def _compute_global_forward_mode(forward_modes):
|
||||||
forward_modes_excluding_idle = [
|
forward_modes_excluding_idle_and_prebuilt = [
|
||||||
x for x in forward_modes if x != ForwardMode.IDLE.value
|
x
|
||||||
|
for x in forward_modes
|
||||||
|
if x != ForwardMode.IDLE.value and x != ForwardMode.PREBUILT.value
|
||||||
]
|
]
|
||||||
|
|
||||||
if not forward_modes_excluding_idle:
|
if not forward_modes_excluding_idle_and_prebuilt:
|
||||||
return ForwardMode.IDLE, False
|
return ForwardMode.IDLE, False
|
||||||
|
|
||||||
forward_mode_agree = TboDPAttentionPreparer._is_all_same(
|
forward_mode_agree = TboDPAttentionPreparer._is_all_same(
|
||||||
forward_modes_excluding_idle
|
forward_modes_excluding_idle_and_prebuilt
|
||||||
)
|
)
|
||||||
|
|
||||||
global_forward_mode = (
|
global_forward_mode = (
|
||||||
ForwardMode(forward_modes_excluding_idle[0]) if forward_mode_agree else None
|
ForwardMode(forward_modes_excluding_idle_and_prebuilt[0])
|
||||||
|
if forward_mode_agree
|
||||||
|
else None
|
||||||
)
|
)
|
||||||
return global_forward_mode, forward_mode_agree
|
return global_forward_mode, forward_mode_agree
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user