minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[bugfix] Support MIXED forward mode in TBO splitter for DP attention (#24241)

Browse Source
This commit is contained in:
Cheng Wan
2026-05-01 16:01:23 -07:00
committed by GitHub
parent 05de73efd1
commit b47fab6f5d
2 changed files with 34 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/two_batch_overlap.py
+2 -2
View File
@@ -81,7 +81,7 @@ def compute_split_seq_index(
extend_lens: Optional[Sequence[int]],
token_num_per_seq: Optional[int],
) -> Optional[int]:
if forward_mode == ForwardMode.EXTEND:
if forward_mode == ForwardMode.EXTEND or forward_mode == ForwardMode.MIXED:
assert extend_lens is not None
return _split_extend_seqs(extend_lens)
elif forward_mode.is_target_verify() or forward_mode.is_decode():
@@ -270,7 +270,7 @@ def compute_split_token_index(
extend_seq_lens: Optional[Sequence[int]],
token_num_per_seq: Optional[int],
) -> int:
if forward_mode == ForwardMode.EXTEND:
if forward_mode == ForwardMode.EXTEND or forward_mode == ForwardMode.MIXED:
assert extend_seq_lens is not None
if _is_two_chunk_split_enabled(extend_seq_lens):
return sum(extend_seq_lens) // 2