diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py index 44b526828..0fdd88ecb 100644 --- a/python/sglang/srt/arg_groups/cuda_graph_hook.py +++ b/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 diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index 992a91b17..88185e363 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/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", diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index e70086e61..05c80805c 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/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).", diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 9b982c232..8fcc17281 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 5f750ee9e..5e2729fd1 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/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.""" diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index b76e513e2..1888f8d7b 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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, - max_seq_len_override=self.MAX_SEQ_LEN_FOR_CAPTURE, + metadata_batch, + max_seq_len_override=max_seq_len, use_prefill_cuda_graph=True, ) assert isinstance(capture_metadata, DSV4Metadata) diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 210323127..4d7f93b6c 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index 94001ba56..64c65856e 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/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, ) diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index fcaf5ef2d..dd2384936 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/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 + ), ) diff --git a/python/sglang/srt/model_executor/cuda_graph_config.py b/python/sglang/srt/model_executor/cuda_graph_config.py index 8a7049532..a61505bae 100644 --- a/python/sglang/srt/model_executor/cuda_graph_config.py +++ b/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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 04b28b8ac..5d832fc97 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 194ca5d37..59d891144 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index baf37a0ec..f980a01a0 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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))) - if get_parallel().tp_rank == 0 - else reversed(self.capture_num_tokens) + tqdm.tqdm(capture_shapes) if get_parallel().tp_rank == 0 else capture_shapes ) 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, diff --git a/test/registered/attention/unittests/dense/test_tbo.py b/test/registered/attention/unittests/dense/test_tbo.py index 27cefebab..90a49fae6 100644 --- a/test/registered/attention/unittests/dense/test_tbo.py +++ b/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__": diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index 03cc90fb5..fc9184160 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/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") - static_forward_batch = SimpleNamespace(name="static") + forward_batch = SimpleNamespace(name="live", max_seq_len_override=None) + static_forward_batch = SimpleNamespace(name="static", max_seq_len_override=None) backend.prepare_forward_metadata_for_breakable_cuda_graph_replay( capture_metadata, diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 1bc6a537d..70da8f1e5 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/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 diff --git a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py index a9f19642a..b76e82bb7 100644 --- a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py +++ b/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 diff --git a/test/registered/unit/managers/scheduler_components/test_dp_attn.py b/test/registered/unit/managers/scheduler_components/test_dp_attn.py index 1cd9230ff..2dbd1cfdd 100644 --- a/test/registered/unit/managers/scheduler_components/test_dp_attn.py +++ b/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, diff --git a/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py b/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py index 99ab09f72..7581b8f7a 100644 --- a/test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py +++ b/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 diff --git a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py index ec33002f4..3ff162b08 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py +++ b/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() diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index f00c9aec6..43ca986e8 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/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 diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index d1f2d9236..3d0202e67 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/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 = ( diff --git a/test/registered/unit/test_server_args_cli_metadata.py b/test/registered/unit/test_server_args_cli_metadata.py index f89b9e068..b0fb4d952 100644 --- a/test/registered/unit/test_server_args_cli_metadata.py +++ b/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( [