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

dp: refactor idle batch logic (#25025)

Browse Source 
Co-authored-by: happierpig <zhaoyilong217@sjtu.edn.cn>
This commit is contained in:
Yilong Zhao
2026-05-26 14:22:25 -07:00
committed by GitHub
co-authored by happierpig
parent dd6f073377
commit 1a05b511e4
1 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/dp_attn.py
+10 -2
View File
@@ -217,14 +217,22 @@ def prepare_mlp_sync_batch_raw(
)
)
need_idle_batch = skip_all_gather or max(mlp_sync_info.global_num_tokens) > 0
# Decide whether to emit idle batch
if skip_all_gather:
# Skip idle batch when attn-dp=1
need_idle_batch = dp_size > 1
else:
need_idle_batch = max(mlp_sync_info.global_num_tokens) > 0
batch_to_gather = local_batch
if need_idle_batch:
batch_to_gather = local_batch
if local_batch is None:
batch_to_gather = local_batch = get_idle_batch()
elif local_batch.forward_mode.is_prebuilt():
# NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
batch_to_gather = local_batch.inner_idle_batch = get_idle_batch()
if batch_to_gather is not None:
_update_gather_batch(
batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
)