[DSV4][BCG] Optimize the heavy memory use of C4 Indexer when BCG is enabled (#36534)

Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>

python/sglang/srt/arg_groups/cuda_graph_hook.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import Any, Optional
 
 from sglang.srt.arg_groups.overrides import (
     attention_backends_of,
@@ -87,6 +87,12 @@ def parse_cuda_graph_config(server_args: Any):
         # decode is implemented; today decode ignores it.
         _set(Phase.DECODE, "tc_compiler", cfg.cuda_graph_tc_compiler)
         _set(Phase.PREFILL, "tc_compiler", cfg.cuda_graph_tc_compiler)
+    if cfg.cuda_graph_prefill_max_context is not None:
+        _set(
+            Phase.PREFILL,
+            "max_context_size",
+            cfg.cuda_graph_prefill_max_context,
+        )
 
     # ---- Explicit JSON config (highest precedence) ----
     for phase, phase_config in explicit_input.items():
@@ -554,6 +560,63 @@ def validate_cuda_graph_config(server_args: Any):
             )
 
 
+def _resolve_max_context_size(
+    *, requested_size: Any, page_size: int, model_context_len: Optional[int]
+) -> int:
+    if requested_size <= 0:
+        raise ValueError("--cuda-graph-prefill-max-context must be a positive integer")
+
+    aligned_size = int((requested_size + page_size - 1) // page_size * page_size)
+    if (
+        model_context_len is not None
+        and model_context_len > 0
+        and aligned_size > model_context_len
+    ):
+        raise ValueError(
+            "--cuda-graph-prefill-max-context exceeds the model context length: "
+            f"aligned size {aligned_size} > {model_context_len}"
+        )
+    if requested_size != aligned_size:
+        logger.info(
+            "Page-aligning prefill CUDA graph max context size %d -> %d "
+            "(page_size=%d).",
+            requested_size,
+            aligned_size,
+            page_size,
+        )
+    return aligned_size
+
+
+def finalize_cuda_graph_prefill_max_context(server_args: Any) -> None:
+    cfg = resolving_view(server_args)
+    requested_size = cfg.cuda_graph_config.prefill.max_context_size
+    if requested_size is None:
+        return
+    page_size = cfg.page_size
+    assert page_size is not None and page_size > 0, (
+        "page_size must be resolved before prefill CUDA graph max context size"
+    )
+
+    max_context_size = _resolve_max_context_size(
+        requested_size=requested_size,
+        page_size=page_size,
+        model_context_len=model_config_of(server_args).context_len,
+    )
+    logger.info(
+        "Prefill CUDA graph max context size: %d; graph keys remain token-only.",
+        max_context_size,
+    )
+    declare_resolution(
+        server_args,
+        "_finalize_cuda_graph_prefill_max_context",
+        cuda_graph_config=with_phase(
+            cfg.cuda_graph_config,
+            Phase.PREFILL,
+            max_context_size=max_context_size,
+        ),
+    )
+
+
 def generate_prefill_cuda_graph_batch_sizes(max_bs: int):
     """
     Generate the list of batch sizes for prefill CUDA graph capture

python/sglang/srt/arg_groups/field_order.py

@@ -230,6 +230,7 @@ POSITIONAL_FIELD_ORDER = (
     "cuda_graph_max_bs_prefill",
     "cuda_graph_bs_decode",
     "cuda_graph_bs_prefill",
+    "cuda_graph_prefill_max_context",
     "cuda_graph_tc_compiler",
     "disable_prefill_cuda_graph",
     "disable_decode_cuda_graph",

python/sglang/srt/arg_groups/fields/exec_.py

@@ -37,6 +37,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
     CudaGraphConfig,
     parse_cuda_graph_config_arg,
 )
+from sglang.srt.utils.common import human_readable_int
 
 
 class ExecFeatures(msgspec.Struct):
@@ -493,6 +494,20 @@ class ExecGraph(msgspec.Struct):
         Optional[List[int]],
         "Explicit list of batch sizes to capture for the prefill cuda graph.",
     ] = None
+    cuda_graph_prefill_max_context: A[
+        Optional[int],
+        Arg(
+            help=(
+                "Maximum context length supported by DeepSeek-V4 breakable/full "
+                "prefill CUDA graphs. Context-shaped attention metadata and "
+                "indexer logits are allocated at this fixed size instead of "
+                "the model maximum. Larger live contexts fall back to eager."
+                f"\n\n{human_readable_int.__doc__}"
+            ),
+            type_parser=human_readable_int,
+            aliases=["--context-bucket"],
+        ),
+    ] = None
     cuda_graph_tc_compiler: A[
         Optional[Literal["eager", "inductor"]],
         "Compiler used by the tc_piecewise backend (currently only the prefill phase consumes it).",

python/sglang/srt/arg_groups/pipeline.py

@@ -190,6 +190,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
         apply_inkling_prefill_cuda_graph_default,
         apply_muse_glimmer_prefill_cuda_graph_max_bs_default,
         disable_prefill_cuda_graph_for_deepseek_trtllm_mla,
+        finalize_cuda_graph_prefill_max_context,
         handle_cuda_graph_config,
     )
 
@@ -369,6 +370,8 @@ def run_resolution_pipeline(server_args: Any) -> None:
     # time; last declarations of the resolution, mirroring that order.
     run_hook(handle_model_capability_adjustments, server_args)
 
+    finalize_cuda_graph_prefill_max_context(server_args)
+
     # Validate after all batch-size declarations are visible.
     run_hook(validate_deepep_v2_speculative_draft, server_args)
     run_hook(validate_deepep_v2_dispatch_token_budget, server_args)

python/sglang/srt/layers/attention/base_attn_backend.py

@@ -153,6 +153,9 @@ class AttentionBackend(ABC):
     # object during capture, and refresh its dynamic fields before each replay.
     use_captured_forward_metadata_for_breakable_cuda_graph: bool = False
 
+    # True when prefill graph metadata can use ForwardBatch.max_seq_len_override.
+    supports_prefill_cuda_graph_max_context_size: bool = False
+
     def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds:
         """Declare where this backend's scheduler-shared reads end per mode.
         Override only for audited deviations from this conservative default."""

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -587,6 +587,7 @@ class DeepseekV4AttnBackend(
     AttentionBackend, C4IndexerBackendMixin, CompressorBackendMixin
 ):
     use_captured_forward_metadata_for_breakable_cuda_graph: bool = True
+    supports_prefill_cuda_graph_max_context_size: bool = True
     supports_ragged_verify_graph: bool = True
     needs_cpu_seq_lens: bool = False
     trtllm_attn: bool = False
@@ -1514,9 +1515,16 @@ class DeepseekV4AttnBackend(
 
         assert self.swa_page_size % SWA_WINDOW == 0 and self.page_size % 128 == 0
         if max_seq_len_override is None:
-            max_seq_len_override = getattr(forward_batch, "max_seq_len_override", None)
+            max_seq_len_override = forward_batch.max_seq_len_override
         if max_seq_len_override is not None:
             max_seq_len = max_seq_len_override
+            if seq_lens_cpu is not None and len(seq_lens_cpu) > 0:
+                actual_max_seq_len = int(seq_lens_cpu.max().item())
+                if actual_max_seq_len > max_seq_len:
+                    raise ValueError(
+                        "Prefill CUDA graph max context size is smaller than the "
+                        f"live context: {max_seq_len=} < {actual_max_seq_len=}"
+                    )
         elif seq_lens_cpu is not None:
             max_seq_len = int(seq_lens_cpu.max().item())
         else:
@@ -1607,9 +1615,10 @@ class DeepseekV4AttnBackend(
     def init_forward_metadata_for_breakable_cuda_graph_capture(
         self, forward_batch: ForwardBatch
     ):
+        max_seq_len = forward_batch.max_seq_len_override or self.MAX_SEQ_LEN_FOR_CAPTURE
         self.forward_metadata = self._build_forward_metadata(
             forward_batch,
-            max_seq_len_override=self.MAX_SEQ_LEN_FOR_CAPTURE,
+            max_seq_len_override=max_seq_len,
             use_prefill_cuda_graph=True,
         )
         return self.forward_metadata
@@ -1624,9 +1633,15 @@ class DeepseekV4AttnBackend(
         # Build graph-compatible metadata against the padded static batch. The
         # batch still carries live seq/extend lens, so the online c128 prefill
         # plan remains batch-specific without constructing a second metadata set.
+        metadata_batch = (
+            static_forward_batch if static_forward_batch is not None else forward_batch
+        )
+        max_seq_len = (
+            metadata_batch.max_seq_len_override or self.MAX_SEQ_LEN_FOR_CAPTURE
+        )
         static_metadata = self._build_forward_metadata(
-            static_forward_batch if static_forward_batch is not None else forward_batch,
+            metadata_batch,
-            max_seq_len_override=self.MAX_SEQ_LEN_FOR_CAPTURE,
+            max_seq_len_override=max_seq_len,
             use_prefill_cuda_graph=True,
         )
         assert isinstance(capture_metadata, DSV4Metadata)

python/sglang/srt/layers/attention/hybrid_attn_backend.py

@@ -93,6 +93,10 @@ class HybridAttnBackend(AttentionBackend):
     def supports_full_cuda_graph_chunked_prefix(self) -> bool:
         return self.prefill_backend.supports_full_cuda_graph_chunked_prefix
 
+    @property
+    def supports_prefill_cuda_graph_max_context_size(self) -> bool:
+        return self.prefill_backend.supports_prefill_cuda_graph_max_context_size
+
     def prepare_full_cuda_graph_chunked_prefix(
         self,
         forward_batch: ForwardBatch,

python/sglang/srt/layers/attention/tbo_backend.py

@@ -28,6 +28,10 @@ class TboAttnBackend(AttentionBackend):
             primary, "extend_dummy_seqs_capped_by_req_pool", False
         )
 
+    @property
+    def supports_prefill_cuda_graph_max_context_size(self) -> bool:
+        return self.primary.supports_prefill_cuda_graph_max_context_size
+
     @classmethod
     def init_new(cls, creator: Callable[[], AttentionBackend]):
         return cls(
@@ -259,4 +263,5 @@ def _build_tbo_child_replay_fb_view(
             else None
         ),
         spec_info=child_spec_info,
+        max_seq_len_override=fb_view.max_seq_len_override,
     )

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -350,6 +350,13 @@ def _local_prefill_cuda_graph_vote(
         capture_hidden_mode=None,
         return_logprob=return_logprob,
         lora_ineligible=prefill_graph_runner.enable_lora,
+        batch_max_context_len=(
+            int(local_batch.seq_lens_cpu.max().item())
+            if prefill_graph_runner.max_context_size is not None
+            and local_batch.seq_lens_cpu is not None
+            and local_batch.seq_lens_cpu.numel() > 0
+            else None
+        ),
     )
 
 

python/sglang/srt/model_executor/cuda_graph_config.py

@@ -72,7 +72,8 @@ ALLOWED_BACKENDS_PER_PHASE = {
 # For prefill, bs carries aggregate-token capture buckets for every backend;
 # full_prefill_max_req separately controls Full's fixed request-slot count.
 # full_prefill_max_req and full_prefill_prefix_chunk_tokens are prefill-only and
-# only meaningful when backend == full.
+# only meaningful when backend == full. max_context_size is shared by the
+# breakable and full prefill body-capture backends.
 ALLOWED_KEYS_PER_PHASE = {
     Phase.DECODE: ("backend", "max_bs", "bs", "tc_compiler"),
     Phase.PREFILL: (
@@ -80,6 +81,7 @@ ALLOWED_KEYS_PER_PHASE = {
         "max_bs",
         "bs",
         "tc_compiler",
+        "max_context_size",
         "full_prefill_max_req",
         "full_prefill_prefix_chunk_tokens",
     ),
@@ -95,6 +97,10 @@ class PhaseConfig:
     bs: Optional[List[int]] = None
     # Only meaningful when backend == tc_piecewise; ignored otherwise.
     tc_compiler: str = "eager"
+    # Effective for both full and breakable backends and currently only DSV4:
+    # fixed maximum context length used by context-shaped prefill graph metadata.
+    # Every token bucket shares this size; larger live contexts run eagerly.
+    max_context_size: Optional[int] = None
     # Only meaningful for the prefill phase with backend == full: max number of
     # request slots baked into each captured graph. Real bs <= full_prefill_max_req
     # reuses the graph (unused slots become zero-length sentinels); larger

python/sglang/srt/model_executor/forward_batch_info.py

@@ -573,6 +573,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     # Preallocated piecewise-graph attention output, set by RadixAttention.
     _attn_output: Optional[torch.Tensor] = None
 
+    # Prefill body-CUDA-graph context limit. Attention backends that allocate
+    # context-shaped metadata use this fixed maximum instead of deriving a
+    # shape from the live batch. None preserves eager/default graph behavior.
+    max_seq_len_override: Optional[int] = None
+
     # For logits and logprobs post processing
     next_token_logits_buffer: torch.Tensor = None
     temperature: torch.Tensor = None

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

@@ -197,6 +197,7 @@ def build_replay_fb_view(
         out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
         out_cache_loc_virtual=forward_batch.out_cache_loc_virtual,
         out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
+        max_seq_len_override=forward_batch.max_seq_len_override,
         # The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE
         # for the backend, which copies the result into its own static buffer and
         # reads THAT in the decode track-save — this slot is never mutated. None

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

@@ -312,6 +312,18 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # --- runner bounds --------------------------------------------
         self.max_num_tokens = max(self.capture_num_tokens)
         self.max_bs = model_runner.req_to_token_pool.size
+        self.max_context_size = prefill_config.max_context_size
+        self._validate_max_context_capacity(
+            max_context_size=self.max_context_size,
+            table_width=model_runner.req_to_token_pool.req_to_token.shape[1],
+        )
+        if (
+            self.prefill_backend_name == Backend.TC_PIECEWISE
+            and self.max_context_size is not None
+        ):
+            # TODO(SYChen123): Plumb max_seq_len_override through TcPiecewise
+            # metadata preparation before enabling the fixed context limit here.
+            self._ignore_max_context_size("tc_piecewise prefill CUDA graph")
 
         # --- capture modes --------------------------------------------
         self.capture_forward_mode = ForwardMode.EXTEND
@@ -434,6 +446,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             max_req = prefill_config.full_prefill_max_req
             assert max_req is not None, "full_prefill_max_req must be resolved"
             self._capture_req_slots = max_req
+
         # BCG/Full record LoRA kernels, so the metadata they read must live in
         # static buffers refreshed in place per batch; unsupported LoRA
         # configs were already routed to the eager runner.
@@ -510,10 +523,22 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             self.backend, BreakableCudaGraphBackend
         ) and should_enable_cp_bcg_capture(server_args)
         if self.enable_cp_bcg_capture:
+            if self.max_context_size is not None:
+                # TODO(SYChen123): Preserve max_seq_len_override through CP's padded
+                # metadata preparation before enabling the fixed context limit.
+                self._ignore_max_context_size("CP breakable prefill CUDA graph")
             self.capture_num_tokens = filter_prefill_cp_bcg_capture_num_tokens(
                 self.capture_num_tokens, server_args
             )
             self.prefill_cp_bcg_input = PrefillCPBCGInput.create(self)
+        if self.max_context_size is not None and not (
+            model_runner.attn_backend.supports_prefill_cuda_graph_max_context_size
+        ):
+            raise ValueError(
+                "--cuda-graph-prefill-max-context is only supported by attention "
+                "backends that implement fixed-context prefill graph metadata; "
+                f"got {type(model_runner.attn_backend).__name__}"
+            )
 
         # Static hidden_states buffer giving the captured graph a stable
         # address; load_batch refreshes it from live spec_info at replay.
@@ -541,7 +566,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             self.use_captured_attn_metadata = model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
         else:
             self.use_captured_attn_metadata = False
-        self.attn_metadata_buffers: Optional[Dict[int, object]] = (
+        self.attn_metadata_buffers: Optional[Dict[ShapeKey, object]] = (
             {} if self.use_captured_attn_metadata else None
         )
 
@@ -862,6 +887,41 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             else table_width
         )
 
+    @staticmethod
+    def _validate_max_context_capacity(
+        *, max_context_size: Optional[int], table_width: int
+    ) -> None:
+        if max_context_size is not None and max_context_size > table_width:
+            raise ValueError(
+                "--cuda-graph-prefill-max-context exceeds the request-to-token "
+                f"pool capacity: requested size {max_context_size} > {table_width}"
+            )
+
+    def _ignore_max_context_size(self, execution_path: str) -> None:
+        if self.max_context_size is None:
+            return
+        logger.warning(
+            "Ignoring prefill CUDA graph max context size %d for %s; this path "
+            "does not yet preserve the fixed metadata extent.",
+            self.max_context_size,
+            execution_path,
+        )
+        self.max_context_size = None
+
+    @staticmethod
+    def _batch_max_context_len(forward_batch: ForwardBatch) -> Optional[int]:
+        seq_lens_cpu = forward_batch.seq_lens_cpu
+        if seq_lens_cpu is not None:
+            if torch.is_tensor(seq_lens_cpu):
+                return int(seq_lens_cpu.max().item()) if seq_lens_cpu.numel() > 0 else 0
+            return max((int(value) for value in seq_lens_cpu), default=0)
+        seq_lens = forward_batch.seq_lens
+        if seq_lens is None or seq_lens.numel() == 0:
+            return None
+        # Prefill normally has seq_lens_cpu. This fallback is for hand-built
+        # batches and intentionally synchronizes rather than guessing a bucket.
+        return int(seq_lens.max().item())
+
     @staticmethod
     def _resolve_prefix_chunk_shape(
         model_runner, capture_req_slots: int
@@ -1054,7 +1114,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         )
 
     def _init_forward_metadata_for_capture(
-        self, forward_batch: ForwardBatch, num_tokens: int
+        self, forward_batch: ForwardBatch, shape_key: ShapeKey
     ) -> None:
         """Capture-time metadata init for the BCG-with-captured-metadata
         contract. For opt-in backends (DSV4), call the BCG-specific entry
@@ -1071,13 +1131,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 )
             )
             assert self.attn_metadata_buffers is not None
-            self.attn_metadata_buffers[num_tokens] = metadata
+            self.attn_metadata_buffers[shape_key] = metadata
 
     def _prepare_forward_metadata_for_replay(
         self,
         forward_batch: ForwardBatch,
         static_forward_batch: ForwardBatch,
-        num_tokens: int,
+        shape_key: ShapeKey,
     ) -> None:
         """Replay-time metadata refresh for the BCG-with-captured-metadata
         contract. For opt-in backends, refresh the stashed per-bucket
@@ -1103,16 +1163,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             padded_view.req_pool_indices = s["req_pool_indices"][:r]
             padded_view.extend_seq_lens = s["extend_seq_lens"][:r]
             padded_view.extend_prefix_lens = s["extend_prefix_lens"][:r]
+            padded_view.max_seq_len_override = static_forward_batch.max_seq_len_override
             attn_backend.init_forward_metadata_out_graph(padded_view)
             return
         if not self.use_captured_attn_metadata:
             attn_backend.init_forward_metadata(forward_batch)
             attn_backend.prepare_prefill_shared_read_snapshot(
-                forward_batch, num_qo_tokens=num_tokens
+                forward_batch, num_qo_tokens=shape_key.size
             )
             return
         assert self.attn_metadata_buffers is not None
-        metadata = self.attn_metadata_buffers[num_tokens]
+        metadata = self.attn_metadata_buffers[shape_key]
         attn_backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
             metadata,
             forward_batch,
@@ -1138,6 +1199,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         capture_hidden_mode,
         return_logprob: bool,
         lora_ineligible: bool = False,
+        batch_max_context_len: Optional[int] = None,
     ) -> bool:
         """Rank-local replay eligibility: the single source of truth for
         ``can_run_graph`` (ForwardBatch, forward time) and the dp mlp-sync
@@ -1180,6 +1242,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             return False
         if return_logprob and not self._uses_eager_prefill_tail():
             return False
+        if self.max_context_size is not None:
+            if (
+                batch_max_context_len is None
+                or batch_max_context_len > self.max_context_size
+            ):
+                return False
         if num_tokens is None:
             return True
         if num_tokens > self.max_num_tokens:
@@ -1206,6 +1274,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             return False
 
         # Non-DP local check (sole decision for tp-only).
+        batch_max_context_len = (
+            self._batch_max_context_len(forward_batch)
+            if self.max_context_size is not None
+            else None
+        )
         if not self.can_replay_locally(
             batch_size=forward_batch.batch_size,
             num_tokens=len(forward_batch.input_ids),
@@ -1222,6 +1295,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     forward_batch
                 )
             ),
+            batch_max_context_len=batch_max_context_len,
         ):
             return False
         if getattr(self, "enable_cp_bcg_capture", False) and is_cp_active(
@@ -1263,7 +1337,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         Returns ``(forward_batch, attn_backend)`` to mirror decode's
         capture_prepare signature.
         """
-        context_length = self.model_runner.model_config.context_len
+        model_context_length = self.model_runner.model_config.context_len
+        context_length = min(
+            self.max_context_size or model_context_length, model_context_length
+        )
         # A prefill bucket is an aggregate token count. Capture it as the
         # fewest synthetic requests, with every request containing no more
         # than context_length tokens.
@@ -1398,6 +1475,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 # refreshes the static batch info with live values.
                 lora_ids=([None] * bs if self._capture_lora else None),
                 return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
+                max_seq_len_override=self.max_context_size,
             )
             self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
         return forward_batch, self.model_runner.attn_backend
@@ -1430,10 +1508,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             self.model_runner.gpu_id,
             empty_cache=False,
         )
+        capture_shapes = list(reversed(self.capture_num_tokens))
+        if self.max_context_size is not None:
+            logger.info(
+                "Capturing %d prefill CUDA graph token shapes with fixed "
+                "max_context_size=%d (before execution variants).",
+                len(capture_shapes),
+                self.max_context_size,
+            )
         capture_range = (
-            tqdm.tqdm(list(reversed(self.capture_num_tokens)))
+            tqdm.tqdm(capture_shapes) if get_parallel().tp_rank == 0 else capture_shapes
-            if get_parallel().tp_rank == 0
-            else reversed(self.capture_num_tokens)
         )
         for num_tokens in capture_range:
             if get_parallel().tp_rank == 0:
@@ -1443,12 +1527,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     empty_cache=False,
                 )
                 capture_range.set_description(
-                    f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
+                    f"Capturing prefill shape ({num_tokens=} {avail_mem=:.2f} GB)"
                 )
             self.capture_one_shape(num_tokens)
             if self._capture_chunked_prefix:
                 for captured_n in self._prefix_capture_variants:
-                    self.capture_one_shape(num_tokens, prefix_num_chunks=captured_n)
+                    self.capture_one_shape(
+                        num_tokens,
+                        prefix_num_chunks=captured_n,
+                    )
 
     def capture_one_shape(self, size: int, *, prefix_num_chunks: int = 0) -> None:
         """Per-shape capture: build dummy ForwardBatch + run_once,
@@ -1496,7 +1583,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             # Reaching into a backend-specific metadata cache here would make
             # this path incompatible with the OSS FlashAttention backend.
         else:
-            self._init_forward_metadata_for_capture(forward_batch, num_tokens)
+            self._init_forward_metadata_for_capture(forward_batch, shape_key)
 
         def run_once():
             # Record LoRA kernels even when capture uses base-model requests.
@@ -1569,6 +1656,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 )
         self.raw_num_tokens = num_tokens
 
+        if self.max_context_size is not None:
+            batch_max_context_len = self._batch_max_context_len(forward_batch)
+            if (
+                batch_max_context_len is None
+                or batch_max_context_len > self.max_context_size
+            ):
+                raise RuntimeError(
+                    "Prefill CUDA graph replay was admitted without a fitting "
+                    "maximum context size"
+                )
+
         bs = forward_batch.batch_size
         self.raw_bs = bs
 
@@ -1707,6 +1805,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 self.capture_return_pooled_hidden_states
                 or forward_batch.return_pooled_hidden_states
             ),
+            max_seq_len_override=self.max_context_size,
         )
         if self._is_full_backend:
             forward_batch.next_token_logits_buffer = (
@@ -1775,8 +1874,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             )
             metadata_forward_batch = static_forward_batch
 
+        shape_key = self._shape_key(static_num_tokens, forward_batch)
         self._prepare_forward_metadata_for_replay(
-            metadata_forward_batch, static_forward_batch, static_num_tokens
+            metadata_forward_batch, static_forward_batch, shape_key
         )
 
         return static_forward_batch
@@ -1860,6 +1960,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         raw_num_tokens: int,
         **kwargs,
     ):
+        assert self.max_context_size is None, (
+            "tc_piecewise replay does not support a fixed prefill context size"
+        )
         with self._prefill_forward_context(
             static_forward_batch,
             num_tokens=static_num_tokens,
@@ -1937,6 +2040,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             )
 
             if self.enable_cp_bcg_capture:
+                assert self.max_context_size is None, (
+                    "CP-v2 BCG replay does not support a fixed prefill context size"
+                )
                 output = execute_prefill_cp_bcg(
                     self,
                     forward_batch,

test/registered/attention/unittests/dense/test_tbo.py

@@ -191,6 +191,7 @@ class TestTboAttnDenseAttentionBackendCorrectness(CustomTestCase):
             encoder_lens=None,
             out_cache_loc=batch.out_cache_loc,
             spec_info=batch.spec_info,
+            max_seq_len_override=700,
         )
 
         # Pure mocks (no `wraps=...`) so the dispatcher's slicing/contract is
@@ -215,6 +216,8 @@ class TestTboAttnDenseAttentionBackendCorrectness(CustomTestCase):
         self.assertEqual(
             child_fbs[1].req_pool_indices.shape[0], capture_bs - split_seq_index
         )
+        self.assertEqual(child_fbs[0].max_seq_len_override, 700)
+        self.assertEqual(child_fbs[1].max_seq_len_override, 700)
 
 
 if __name__ == "__main__":

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -579,8 +579,8 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
             return replay_metadata
 
         backend._build_forward_metadata = fake_build_forward_metadata
-        forward_batch = SimpleNamespace(name="live")
+        forward_batch = SimpleNamespace(name="live", max_seq_len_override=None)
-        static_forward_batch = SimpleNamespace(name="static")
+        static_forward_batch = SimpleNamespace(name="static", max_seq_len_override=None)
 
         backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
             capture_metadata,

test/registered/cp/test_cp_strategy_unit.py

@@ -135,6 +135,7 @@ class TestPrefillCPBCGReplay(CustomTestCase):
         runner.has_mha_companion_layers = False
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.capture_num_tokens = [2048, 2304]
+        runner.max_context_size = None
         runner.max_num_tokens = 2304
         runner.enable_cp_bcg_capture = True
         return runner

test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py

@@ -43,6 +43,7 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase):
         runner.has_mha_companion_layers = backend == Backend.BREAKABLE
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.capture_num_tokens = [4, 16]
+        runner.max_context_size = None
         runner.max_num_tokens = 16
         return runner
 

test/registered/unit/managers/scheduler_components/test_dp_attn.py

@@ -84,6 +84,7 @@ class TestDecodeToExtendConversionVote(CustomTestCase):
     def _vote(self, *, beam):
         runner = Mock(spec=dp_attn.PrefillCudaGraphRunner)
         runner.enable_lora = False
+        runner.max_context_size = None
         runner.can_replay_locally.return_value = True
         batch = SimpleNamespace(
             forward_mode=ForwardMode.DECODE,

test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py

@@ -73,6 +73,7 @@ class TestHiddenStateGraphRecapture(CustomTestCase):
         runner._capture_chunked_prefix = False
         runner.capture_hidden_mode = capture_hidden_mode
         runner.capture_num_tokens = [4]
+        runner.max_context_size = None
         runner.max_num_tokens = 4
         return runner
 

test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py

@@ -2,6 +2,9 @@ import unittest
 from types import SimpleNamespace
 from unittest import mock
 
+import torch
+
+import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
 from sglang.srt.layers.moe.utils import MoeA2ABackend
 from sglang.srt.model_executor import forward_batch_info
 from sglang.srt.model_executor.cuda_graph_config import Backend
@@ -13,6 +16,7 @@ from sglang.srt.model_executor.forward_batch_info import (
 from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
     PrefillCudaGraphRunner,
 )
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -29,6 +33,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
         runner.has_mha_companion_layers = False
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.capture_num_tokens = [4, 16]
+        runner.max_context_size = None
         runner.max_num_tokens = 16
         return runner
 
@@ -44,6 +49,8 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
             return_logprob=False,
             input_ids=list(range(num_tokens)),
             extend_prefix_lens_cpu=[0],
+            seq_lens_cpu=torch.tensor([num_tokens], dtype=torch.int64),
+            seq_lens=torch.tensor([num_tokens], dtype=torch.int64),
         )
 
     def test_rejects_more_than_two_x_token_padding(self):
@@ -67,7 +74,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
         runner._prepare_forward_metadata_for_replay(
             forward_batch,
             static_forward_batch,
-            num_tokens=16,
+            shape_key=ShapeKey(size=16),
         )
 
         attn_backend.init_forward_metadata.assert_called_once_with(forward_batch)
@@ -137,6 +144,28 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
                 )
             )
 
+    def test_rejects_context_above_fixed_maximum(self):
+        runner = self._make_runner()
+        runner.max_context_size = 700
+
+        much_shorter = self._make_forward_batch(4)
+        much_shorter.seq_lens_cpu.fill_(200)
+        self.assertTrue(runner.can_run_graph(much_shorter))
+
+        uncovered = self._make_forward_batch(4)
+        uncovered.seq_lens_cpu.fill_(701)
+        self.assertFalse(runner.can_run_graph(uncovered))
+
+    def test_unsupported_path_ignores_max_context_size(self):
+        runner = self._make_runner()
+        runner.max_context_size = 1024
+
+        with self.assertLogs(runner_module.logger, level="WARNING") as logs:
+            runner._ignore_max_context_size("test path")
+
+        self.assertIsNone(runner.max_context_size)
+        self.assertIn("fixed metadata extent", "\n".join(logs.output))
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -244,6 +244,8 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
     def test_static_batch_preserves_consumed_multimodal_embeddings(self):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner.capture_num_tokens = [4]
+        runner.max_context_size = None
+        runner._capture_chunked_prefix = False
         runner.buffer_registry = _FakeBatchRegistry()
         runner.model_runner = SimpleNamespace(attn_tp_sequence_sharded=lambda _: False)
         runner.enable_cp_bcg_capture = False
@@ -477,6 +479,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         runner.capture_hidden_mode = CaptureHiddenMode.NULL
         runner.max_num_tokens = 32
         runner.capture_num_tokens = [4]
+        runner.max_context_size = None
         runner.backend = SimpleNamespace()
         runner.prefill_backend_name = Backend.FULL
         runner.has_mha_companion_layers = False

test/registered/unit/server_args/test_server_args.py

@@ -21,6 +21,7 @@ from sglang.srt.arg_groups.attention_hook import (
 from sglang.srt.arg_groups.cuda_graph_hook import (
     apply_cuda_graph_compatibility,
     disable_tc_piecewise_cudagraph_if_incompatible,
+    finalize_cuda_graph_prefill_max_context,
     handle_cuda_graph_config,
 )
 from sglang.srt.arg_groups.hicache_hook import (
@@ -2114,6 +2115,35 @@ class TestCudaGraphConfigDataclassAccess(CustomTestCase):
         self.assertEqual(config.compiler, "eager")
 
 
+class TestCudaGraphPrefillMaxContextResolution(CustomTestCase):
+    @staticmethod
+    def _make_args(max_context_size, model_context_len=4096, page_size=64):
+        args = ServerArgs(
+            model_path="dummy",
+            page_size=page_size,
+            cuda_graph_config=CudaGraphConfig(
+                prefill=PhaseConfig(
+                    backend=Backend.BREAKABLE,
+                    max_context_size=max_context_size,
+                )
+            ),
+        )
+        args._model_config = SimpleNamespace(context_len=model_context_len)
+        return args
+
+    def test_rejects_invalid_values_during_resolution(self):
+        cases = (
+            (0, "positive integer"),
+            (-1, "positive integer"),
+            (4097, "model context length"),
+        )
+        for max_context_size, expected_error in cases:
+            with self.subTest(max_context_size=max_context_size):
+                args = self._make_args(max_context_size)
+                with self.assertRaisesRegex(ValueError, expected_error):
+                    finalize_cuda_graph_prefill_max_context(args)
+
+
 class TestPipelineParallelPrefillCudaGraphPolicy(CustomTestCase):
     def test_pp_prefill_graph_is_opt_in(self):
         cases = (

test/registered/unit/test_server_args_cli_metadata.py

@@ -45,6 +45,11 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
             self.actions_by_option["--prefill-delayer-forward-passes-buckets"].nargs,
             "+",
         )
+        self.assertIs(
+            self.actions_by_option["--cuda-graph-prefill-max-context"].type,
+            human_readable_int,
+        )
+        self.assertIsNone(self.actions_by_option["--context-bucket"].nargs)
         self.assertEqual(
             self.actions_by_option["--schedule-policy"].choices,
             [
@@ -76,6 +81,19 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
                 self.assertEqual(args.dp_size, 3)
                 self.assertEqual(ServerArgs.from_cli_args(args).dp_size, 3)
 
+    def test_prefill_max_context_accepts_human_readable_values(self):
+        for option in (
+            "--cuda-graph-prefill-max-context",
+            "--context-bucket",
+        ):
+            with self.subTest(option=option):
+                args = self.parser.parse_args(["--model", "dummy", option, "200k"])
+
+                self.assertEqual(
+                    ServerArgs.from_cli_args(args).cuda_graph_prefill_max_context,
+                    200_000,
+                )
+
     def test_migrated_and_manual_options_parse_together(self):
         args = self.parser.parse_args(
             [