EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

support idle batch for TBO (#8233)

Browse Source
This commit is contained in:
ZhichenJiang
2025-07-25 11:43:52 +08:00
committed by GitHub
parent d40846d456
commit f4674df646
1 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/two_batch_overlap.py
View File

@@ -341,15 +341,18 @@ class TboDPAttentionPreparer:
@staticmethod @staticmethod
def _compute_global_forward_mode(forward_modes): def _compute_global_forward_mode(forward_modes):
converted_forward_modes = [ forward_modes_excluding_idle = [
ForwardMode.DECODE.value if x == ForwardMode.IDLE.value else x x for x in forward_modes if x != ForwardMode.IDLE.value
for x in forward_modes
] ]
if not forward_modes_excluding_idle:
return ForwardMode.IDLE, False
forward_mode_agree = TboDPAttentionPreparer._is_all_same( forward_mode_agree = TboDPAttentionPreparer._is_all_same(
converted_forward_modes forward_modes_excluding_idle
) )
global_forward_mode = ( 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 return global_forward_mode, forward_mode_agree