Fix prefill CP graph overflow with larger bucket search (#33906)

python/sglang/srt/layers/cp/bcg.py

@@ -17,16 +17,19 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch
 
+from sglang.srt.layers.cp.base import get_cp_strategy
+from sglang.srt.layers.cp.padding import get_cp_padding_align_size
 from sglang.srt.layers.cp.utils import (
     cp_gather_after_forward,
     cp_split_before_forward,
     enable_cp_v2,
     prepare_cp_forward,
 )
+from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
 
 if TYPE_CHECKING:
@@ -106,6 +109,67 @@ class PrefillCPBCGInput:
                 ),
             )
 
+    def required_local_tokens(self, extend_seq_lens: Any) -> Optional[int]:
+        """Return the aligned CP-local rows required by a live zigzag layout."""
+        strategy = get_cp_strategy()
+        if not isinstance(strategy, ZigzagCPStrategy) or extend_seq_lens is None:
+            return None
+
+        cp_segment_num = strategy.cp_size * 2
+        per_rank_logical_tokens = [0] * strategy.cp_size
+        for raw_length in extend_seq_lens:
+            base, remainder = divmod(int(raw_length), cp_segment_num)
+            for rank in range(strategy.cp_size):
+                opposite_rank = cp_segment_num - 1 - rank
+                per_rank_logical_tokens[rank] += (
+                    base * 2 + int(rank < remainder) + int(opposite_rank < remainder)
+                )
+        align_size = get_cp_padding_align_size()
+        return (
+            (max(per_rank_logical_tokens) + align_size - 1) // align_size * align_size
+        )
+
+    def select_replay_bucket(
+        self,
+        *,
+        num_tokens: int,
+        required_local_tokens: int,
+        capture_num_tokens: list[int],
+        max_padding_factor: int,
+    ) -> Optional[int]:
+        """Return the smallest global capture whose CP-local rows fit."""
+        max_num_tokens = num_tokens * max_padding_factor
+        for bucket in capture_num_tokens:
+            if bucket < num_tokens:
+                continue
+            if bucket > max_num_tokens:
+                break
+            captured_local_tokens = self.bucket_local_tokens.get(bucket)
+            if (
+                captured_local_tokens is not None
+                and required_local_tokens <= captured_local_tokens
+            ):
+                return bucket
+        return None
+
+    def select_replay_bucket_for_batch(
+        self,
+        *,
+        num_tokens: int,
+        extend_seq_lens: Any,
+        capture_num_tokens: list[int],
+        max_padding_factor: int,
+    ) -> Optional[int]:
+        required_local_tokens = self.required_local_tokens(extend_seq_lens)
+        if required_local_tokens is None:
+            return None
+        return self.select_replay_bucket(
+            num_tokens=num_tokens,
+            required_local_tokens=required_local_tokens,
+            capture_num_tokens=capture_num_tokens,
+            max_padding_factor=max_padding_factor,
+        )
+
     def prepare(
         self,
         runner: PrefillCudaGraphRunner,

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -63,6 +63,7 @@ from sglang.srt.layers.cp.bcg import (
     execute_prefill_cp_bcg,
     filter_prefill_cp_bcg_capture_num_tokens,
 )
+from sglang.srt.layers.cp.utils import is_cp_v2_active
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
     set_dp_buffer_len,
@@ -1137,6 +1138,20 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             ),
         ):
             return False
+        if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(
+            forward_batch
+        ):
+            assert self.prefill_cp_bcg_input is not None
+            if (
+                self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
+                    num_tokens=len(forward_batch.input_ids),
+                    extend_seq_lens=forward_batch.extend_seq_lens_cpu,
+                    capture_num_tokens=self.capture_num_tokens,
+                    max_padding_factor=_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR,
+                )
+                is None
+            ):
+                return False
         # Multi-req replay is supported by body-capture backends via the
         # layer_model.forward monkey-patch in replay(): the captured graph runs
         # the transformer stack, then the outer model.forward runs
@@ -1418,6 +1433,22 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         """
         num_tokens = len(forward_batch.input_ids)
         static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
+        if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(
+            forward_batch
+        ):
+            assert self.prefill_cp_bcg_input is not None
+            static_num_tokens = (
+                self.prefill_cp_bcg_input.select_replay_bucket_for_batch(
+                    num_tokens=num_tokens,
+                    extend_seq_lens=forward_batch.extend_seq_lens_cpu,
+                    capture_num_tokens=self.capture_num_tokens,
+                    max_padding_factor=_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR,
+                )
+            )
+            if static_num_tokens is None:
+                raise RuntimeError(
+                    "Prefill CUDA graph replay was admitted without a fitting bucket"
+                )
         self.raw_num_tokens = num_tokens
 
         bs = forward_batch.batch_size