From f4674df646ca8a5515dfdc93677f7bdc052416c6 Mon Sep 17 00:00:00 2001 From: ZhichenJiang <1147802470@qq.com> Date: Fri, 25 Jul 2025 11:43:52 +0800 Subject: [PATCH] support idle batch for TBO (#8233) --- python/sglang/srt/two_batch_overlap.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/two_batch_overlap.py b/python/sglang/srt/two_batch_overlap.py index 3fdf2a1f7..74bc1ba85 100644 --- a/python/sglang/srt/two_batch_overlap.py +++ b/python/sglang/srt/two_batch_overlap.py @@ -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