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