[CPU] Avoid prefill CP predicates during decode graph capture (#39690)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
2026-09-18 09:21:05 +08:00
parent f65c70bb7d
commit a407915c17
1 changed files with 4 additions and 1 deletions
@@ -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