From 7eb27372b34679cdc42d22875b328b103b1ef11d Mon Sep 17 00:00:00 2001 From: "Po-Han Huang (NVIDIA)" <53919306+nvpohanh@users.noreply.github.com> Date: Tue, 4 Aug 2026 07:56:59 +0800 Subject: [PATCH] fix(server): capture legal multi-request prefill CUDA graph batches (#30206) --- .../srt/model_executor/cuda_graph_config.py | 13 +++--- .../cuda_graph_setup.py | 43 ++++++++++++++++- .../runner/prefill_cuda_graph_runner.py | 46 +++++++++++++------ python/sglang/srt/server_args.py | 8 ---- 4 files changed, 79 insertions(+), 31 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_config.py b/python/sglang/srt/model_executor/cuda_graph_config.py index 37c07a4a4..04db53ce7 100644 --- a/python/sglang/srt/model_executor/cuda_graph_config.py +++ b/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 - # captured bucket. Opt-in: the padding waste is the operator's call. + # bucket with a fixed request-slot count; replay pads num_tokens up to + # 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 -# tc_piecewise, request count for breakable) — one shape knob 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. 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 - # (bs-invariant via torch.compile). None auto-derives chunked_prefill_size // 512. + # batches fall back to eager. Ignored by BCG and TC_PIECEWISE. None + # 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 diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 6d7a0160a..f4ded560b 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/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}, " 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 c334f7f27..63f477fc3 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 @@ -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: - # Auto: scale request slots with the chunked prefill size. - max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1) - self._capture_req_slots = min(max_req, self.max_bs) + 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. @@ -1152,19 +1150,37 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): Returns ``(forward_batch, attn_backend)`` to mirror decode's capture_prepare signature. """ - bs = self._capture_req_slots - # Slot 0 carries num_tokens; slots 1..bs-1 are zero-length sentinels. - lens_cpu = [num_tokens] + [0] * (bs - 1) - start_loc_cpu = [0] + [num_tokens] * (bs - 1) + context_length = self.model_runner.model_config.context_len + # 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. + 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), - "orig_seq_lens": torch.tensor(lens_cpu, device=self.device), - "extend_seq_lens": torch.tensor(lens_cpu, device=self.device), + "seq_lens": torch.tensor(capture_seq_lens, device=self.device), + "orig_seq_lens": torch.tensor(capture_seq_lens, 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_logprob_start_lens_cpu=list(lens_cpu), + extend_seq_lens_cpu=list(capture_seq_lens), + 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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 32aba514f..9c2f3895b 100644 --- a/python/sglang/srt/server_args.py +++ b/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(