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

Tiny simplify can_run_dp_cuda_graph gather logic (#12891)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-09 11:15:36 +08:00
committed by GitHub
parent 5236290399
commit 49543be9da
1 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
+3 -5
View File
@@ -100,10 +100,9 @@ def prepare_mlp_sync_batch_raw(
# When return_logprob = False, only need last token per request # When return_logprob = False, only need last token per request
num_tokens_for_logprob = local_batch.batch_size() num_tokens_for_logprob = local_batch.batch_size()
if local_batch is None or local_batch.forward_mode.is_decode_or_idle(): can_cuda_graph = (
can_cuda_graph = 1 local_batch is None or local_batch.forward_mode.is_decode_or_idle()
else: ) and not disable_cuda_graph
can_cuda_graph = 0
is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
@@ -151,7 +150,6 @@ def prepare_mlp_sync_batch_raw(
local_batch.global_forward_mode = global_forward_mode local_batch.global_forward_mode = global_forward_mode
# Check forward mode for cuda graph # Check forward mode for cuda graph
if not disable_cuda_graph:
local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
return local_batch return local_batch