minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Po-Han Huang (NVIDIA)
2026-08-03 16:56:59 -07:00
committed by GitHub
parent b819d2fb5b
commit 7eb27372b3
4 changed files with 79 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/cuda_graph_config.py
+7 -6
View File
@@ -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
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+41 -2
View File
@@ -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
+31 -15
View File
@@ -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,
python/sglang/srt/server_args.py
-8
View File
@@ -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(