fix(server): capture legal multi-request prefill CUDA graph batches (#30206)

python/sglang/srt/model_executor/cuda_graph_config.py

@@ -53,8 +53,9 @@ ALLOWED_BACKENDS_PER_PHASE = {
         Backend.DISABLED,
     ),
     # full for prefill captures one whole-forward graph per num_tokens
-    # bucket (bs=1 only); replay pads num_tokens up to the nearest
+    # bucket with a fixed request-slot count; replay pads num_tokens up to
-    # captured bucket. Opt-in: the padding waste is the operator's call.
+    # the nearest captured bucket. Opt-in: the padding waste is the
+    # operator's call.
     Phase.PREFILL: (
         Backend.FULL,
         Backend.BREAKABLE,
@@ -66,8 +67,8 @@ ALLOWED_BACKENDS_PER_PHASE = {
 # Per-phase settings schema. Keys other than backend are runner-level
 # (read by any backend in that phase); tc_compiler is the lone
 # backend-specific knob (only meaningful when backend == tc_piecewise).
-# For prefill, bs carries the captured shape size (token count for Full and
+# For prefill, bs carries aggregate-token capture buckets for every backend;
-# tc_piecewise, request count for breakable) — one shape knob per phase.
+# 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.
 ALLOWED_KEYS_PER_PHASE = {
@@ -95,8 +96,8 @@ class PhaseConfig:
     # 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
-    # batches fall back to eager. Ignored by BCG (bs=1 only) and TC_PIECEWISE
+    # batches fall back to eager. Ignored by BCG and TC_PIECEWISE. None
-    # (bs-invariant via torch.compile). None auto-derives chunked_prefill_size // 512.
+    # auto-derives chunked_prefill_size // 512.
     full_prefill_max_req: Optional[int] = None
     # Only meaningful for Full prefill CUDA graphs that capture a distinct
     # cached-prefix topology: aggregate cached-prefix tokens represented by one

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -211,6 +211,47 @@ def capture_prefill_graph(
         logger.warning("Disable prefill CUDA graph because the capture size is not set")
         return None
 
+    prefill_config = model_runner.server_args.cuda_graph_config.prefill
+    prefill_backend = prefill_config.backend
+    context_length = model_runner.model_config.context_len
+    if prefill_backend == Backend.FULL:
+        max_capture_requests = prefill_config.full_prefill_max_req
+        if max_capture_requests is None:
+            max_capture_requests = max(
+                model_runner.server_args.chunked_prefill_size // 512, 1
+            )
+        max_capture_requests = min(
+            max_capture_requests, model_runner.req_to_token_pool.size
+        )
+        # Resolve Full's fixed request-axis shape once, just like bs below.
+        prefill_config.full_prefill_max_req = max_capture_requests
+    else:
+        max_capture_requests = model_runner.req_to_token_pool.size
+    # The capture dummy batch has at most max_capture_requests rows, and
+    # each row can contain at most context_length tokens. Their product is
+    # therefore the largest aggregate-token bucket capture can represent.
+    max_capture_tokens = max_capture_requests * context_length
+    capture_num_tokens = sorted(
+        num_tokens
+        for num_tokens in prefill_config.bs
+        if num_tokens <= max_capture_tokens
+    )
+    # Resolve the context- and request-capacity-bounded buckets once before
+    # constructing the runner so every backend consumes the same config.
+    prefill_config.bs = capture_num_tokens
+    if not capture_num_tokens:
+        logger.warning(
+            "Disable prefill CUDA graph capture because no configured "
+            "capture size fits backend=%s with max_capture_tokens=%s "
+            "(max_capture_requests=%s, context_length=%s, request-pool size=%s).",
+            prefill_backend,
+            max_capture_tokens,
+            max_capture_requests,
+            context_length,
+            model_runner.req_to_token_pool.size,
+        )
+        return eager_runner
+
     # Collect attention layers and moe layers from the model. Keep a VLM
     # wrapper that exposes ``language_model`` unchanged: assigning it to
     # ``model`` would register a duplicate module alias and duplicate the
@@ -251,7 +292,6 @@ def capture_prefill_graph(
 
     tic = time.perf_counter()
     before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
-    prefill_backend = model_runner.server_args.cuda_graph_config.prefill.backend
     if should_skip_auto_prefill_cuda_graph_for_memory(
         before_mem,
         getattr(model_runner.server_args, "_cuda_graph_config_locked", set()),
@@ -268,7 +308,6 @@ def capture_prefill_graph(
 
     role = "draft" if model_runner.is_draft_worker else "target"
     capture_name = f"{role} prefill"
-    capture_num_tokens = sorted(model_runner.server_args.cuda_graph_config.prefill.bs)
     logger.info(
         f"Capture {capture_name} CUDA graph begin. "
         f"backend={prefill_backend}, num_tokens={capture_num_tokens}, "

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

@@ -377,10 +377,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self._is_full_backend = isinstance(self.backend, FullCudaGraphBackend)
         if self._is_full_backend:
             max_req = prefill_config.full_prefill_max_req
-            if max_req is None:
+            assert max_req is not None, "full_prefill_max_req must be resolved"
-                # Auto: scale request slots with the chunked prefill size.
+            self._capture_req_slots = max_req
-                max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1)
-            self._capture_req_slots = min(max_req, self.max_bs)
         # 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.
@@ -1152,19 +1150,37 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         Returns ``(forward_batch, attn_backend)`` to mirror decode's
         capture_prepare signature.
         """
-        bs = self._capture_req_slots
+        context_length = self.model_runner.model_config.context_len
-        # Slot 0 carries num_tokens; slots 1..bs-1 are zero-length sentinels.
+        # A prefill bucket is an aggregate token count. Capture it as the
-        lens_cpu = [num_tokens] + [0] * (bs - 1)
+        # fewest synthetic requests, with every request containing no more
-        start_loc_cpu = [0] + [num_tokens] * (bs - 1)
+        # than context_length tokens.
+        capture_seq_lens = [
+            min(context_length, num_tokens - start)
+            for start in range(0, num_tokens, context_length)
+        ]
+        if self.prefill_backend_name == Backend.FULL:
+            # Full captures a fixed request-axis shape; unused slots are
+            # zero-length sentinels after the context-bounded real requests.
+            capture_seq_lens += [0] * (self._capture_req_slots - len(capture_seq_lens))
+        bs = len(capture_seq_lens)
+        assert bs <= self.max_bs, (
+            f"Capture batch needs {bs} request slots for {num_tokens} tokens, "
+            f"but the request pool has only {self.max_bs}."
+        )
+        capture_start_locs = [0]
+        for seq_len in capture_seq_lens[:-1]:
+            capture_start_locs.append(capture_start_locs[-1] + seq_len)
 
         with torch.device(self.device):
             shape_inputs = {
                 "req_pool_indices": torch.arange(bs, device=self.device),
-                "seq_lens": torch.tensor(lens_cpu, device=self.device),
+                "seq_lens": torch.tensor(capture_seq_lens, device=self.device),
-                "orig_seq_lens": torch.tensor(lens_cpu, device=self.device),
+                "orig_seq_lens": torch.tensor(capture_seq_lens, device=self.device),
-                "extend_seq_lens": torch.tensor(lens_cpu, device=self.device),
+                "extend_seq_lens": torch.tensor(capture_seq_lens, device=self.device),
                 "extend_prefix_lens": torch.zeros((bs,), dtype=torch.int64),
-                "extend_start_loc": torch.tensor(start_loc_cpu, device=self.device),
+                "extend_start_loc": torch.tensor(
+                    capture_start_locs, device=self.device
+                ),
             }
         if self._prefill_static_buffers is not None:
             s = self._prefill_static_buffers
@@ -1215,7 +1231,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 seq_lens=shape_inputs["seq_lens"],
                 next_token_logits_buffer=self._next_token_logits_buffer(bs),
                 orig_seq_lens=shape_inputs["orig_seq_lens"],
-                seq_lens_cpu=torch.tensor(lens_cpu, device="cpu"),
+                seq_lens_cpu=torch.tensor(capture_seq_lens, device="cpu"),
                 out_cache_loc=_slot("out_cache_loc"),
                 seq_lens_sum=num_tokens,
                 mamba_track_indices=(
@@ -1240,8 +1256,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                 extend_prefix_lens=shape_inputs["extend_prefix_lens"],
                 extend_start_loc=shape_inputs["extend_start_loc"],
                 extend_prefix_lens_cpu=[0] * bs,
-                extend_seq_lens_cpu=list(lens_cpu),
+                extend_seq_lens_cpu=list(capture_seq_lens),
-                extend_logprob_start_lens_cpu=list(lens_cpu),
+                extend_logprob_start_lens_cpu=list(capture_seq_lens),
                 positions=_slot("positions"),
                 global_num_tokens_gpu=global_num_tokens_gpu,
                 global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,

python/sglang/srt/server_args.py

@@ -4740,14 +4740,6 @@ class ServerArgs:
                     prefill_cuda_graph_config.max_bs, 4096
                 )
 
-        # Clamp to context_length if explicitly set — prevents prefill CG
-        # warmup from compiling graphs with more tokens than the model
-        # buffers can hold, which causes illegal memory access (#21112).
-        if self.context_length is not None:
-            prefill_cuda_graph_config.max_bs = min(
-                prefill_cuda_graph_config.max_bs, self.context_length
-            )
-
         if prefill_cuda_graph_config.bs is None:
             prefill_cuda_graph_config.bs = (
                 self._generate_prefill_cuda_graph_batch_sizes(