From a407915c175e6a3b22ac1d4941592b3f7026b265 Mon Sep 17 00:00:00 2001 From: YanbingJiang Date: Fri, 18 Sep 2026 09:21:05 +0800 Subject: [PATCH] [CPU] Avoid prefill CP predicates during decode graph capture (#39690) Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/communicator.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 323d4c309..05af063fb 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -915,7 +915,10 @@ class LayerCommunicator: return True if forward_batch.dp_padding_mode.is_max_len(): return True - if dsa_use_prefill_cp(forward_batch) or is_mla_cp_active(forward_batch): + # Prefill CP predicates must stay out of decode graph capture. + if forward_batch.forward_mode.is_context_parallel_extend() and ( + dsa_use_prefill_cp(forward_batch) or is_mla_cp_active(forward_batch) + ): return True if get_attn_tp_context().input_scattered and not self.is_last_layer: return True