[Experimental] Full Cuda Graph Support for Prefill by Oasis-Git · Pull Request #27988 · sgl-project/sglang

Oasis-Git · 2026-06-12T02:43:24Z

Adds "full" as a prefill cuda-graph backend (previously decode-only): one whole-forward torch.cuda.CUDAGraph per num_tokens bucket, reusing the existing FullCudaGraphBackend unchanged. Targets small prefills on fast GPUs where per-forward launch/python overhead dominates.

Design:

Token axis: replay pads num_tokens up to the nearest captured bucket (same ladder as the other prefill backends).
Request axis: each graph is captured with a fixed slot count (FULL_CG_PREFILL_REQ_SLOTS=16); real batches with bs <= slots reuse the graph by padding the metadata with zero-length sentinel requests (zeroed every replay), so one graph serves bs in [1, 16]. Larger batches fall back to eager via can_run.
Attention metadata follows the decode-style 2-step contract (init_forward_metadata_out_graph before capture/replay). Implemented for flashinfer (dedicated use_cuda_graph wrappers, one set shared across buckets, 2GB dedicated workspace; split-kv must stay enabled in cudagraph mode -- its block_valid_mask is what lets padded/stale tiles of the fixed captured grid exit early, measured 165x) and for the FlashAttention backend (fa4; dedicated buffers, cu_seqlens_q never aliased to cu_seqlens_k, max_seq_len_q/k baked as upper bounds; page_size=1 only for now). Other attention backends fail loudly at capture during startup.
disable_split_kv threaded through FlashInferIndicesUpdaterPrefill for parity with the decode updater.

Results (Qwen3-8B, B200, mgsm_en 1319 examples, threshold 0.80): flashinfer 0.836-0.84, fa4 0.852; no default-config regression (0.848). Small-prefill forward (30 tok): fa4 4.07 ms / flashinfer 4.41 ms vs BCG 7.35 ms (1.7x), identical GPU busy -- the win is launch overhead.

Test: test/registered/cuda_graph/full_prefill/ mirrors the breakable integration test with --cuda-graph-backend-prefill=full.

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ✅ Run #28691874062
Latest PR Test (Extra): ✅ Run #29110980288

gemini-code-assist · 2026-06-12T02:43:27Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Adds "full" as a prefill cuda-graph backend (previously decode-only): one whole-forward torch.cuda.CUDAGraph per num_tokens bucket, reusing the existing FullCudaGraphBackend unchanged. Targets small prefills on fast GPUs where per-forward launch/python overhead dominates. Design: - Token axis: replay pads num_tokens up to the nearest captured bucket (same ladder as the other prefill backends). - Request axis: each graph is captured with a fixed slot count (FULL_CG_PREFILL_REQ_SLOTS=16); real batches with bs <= slots reuse the graph by padding the metadata with zero-length sentinel requests (zeroed every replay), so one graph serves bs in [1, 16]. Larger batches fall back to eager via can_run. - Attention metadata follows the decode-style 2-step contract (init_forward_metadata_out_graph before capture/replay). Implemented for flashinfer (dedicated use_cuda_graph wrappers, one set shared across buckets, 2GB dedicated workspace; split-kv must stay enabled in cudagraph mode -- its block_valid_mask is what lets padded/stale tiles of the fixed captured grid exit early, measured 165x) and for the FlashAttention backend (fa4; dedicated buffers, cu_seqlens_q never aliased to cu_seqlens_k, max_seq_len_q/k baked as upper bounds; page_size=1 only for now). Other attention backends fail loudly at capture during startup. - disable_split_kv threaded through FlashInferIndicesUpdaterPrefill for parity with the decode updater. Results (Qwen3-8B, B200, mgsm_en 1319 examples, threshold 0.80): flashinfer 0.836-0.84, fa4 0.852; no default-config regression (0.848). Small-prefill forward (30 tok): fa4 4.07 ms / flashinfer 4.41 ms vs BCG 7.35 ms (1.7x), identical GPU busy -- the win is launch overhead. Test: test/registered/cuda_graph/full_prefill/ mirrors the breakable integration test with --cuda-graph-backend-prefill=full. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Previously, Full prefill CG captured model.forward including the LM head and logits_processor. The output buffer _copy_logits_to_buffer wrote into was sized (FULL_CG_PREFILL_REQ_SLOTS, vocab) per captured bucket. With 58 buckets and Qwen3-8B vocab=151,936, that's ~9.7 MB per bucket at R=16 (563 MB total in graph pool); at R=1024 it explodes to ~622 MB per bucket (36 GB total) — OOM on first capture. BCG already faces the same shape problem at the request axis and solves it by capturing only layer_model.forward (the transformer body) and running the outer model.forward tail (LM head + logits_processor) eagerly on top with live multi-req metadata. The runner already plumbs this: self.layer_model dispatches in _run_forward (capture) and in replay's monkey-patch path. The Full backend just wasn't enabled for it. Three small changes: 1. Extend the layer_model gate at construction to include Full. 2. In replay's patched layer_model.forward, slice the static hidden_states down to raw_num_tokens for Full (BCG bs=1 has no sentinels so the slice would be a no-op). 3. In the outer model.forward call after monkey-patching, use the raw user-facing forward_batch (not the R-padded static_forward_batch) for Full so the eager tail's lm_head + logits_processor compute at raw_bs / raw_num_tokens — bounded by real load, not by R. Validated on H200, Qwen3-8B, mgsm_en (cg-refactor/cg-prefill): R=16 + fa3 : 0.848 (baseline 0.840), CG cost 1.90 GB (+240 MB) R=16 + flashinfer : 0.824 (baseline 0.852), CG cost 3.35 GB (+180 MB) R=1024 + fa3 : 0.832, CG cost 2.04 GB (+140 MB over R=16) R=1024 + fa3 OOM'd before this change (logits buffer 622 MB per bucket). After: page_table is now the only R-dependent allocation (167 MB at R=1024, vs negligible at R=16), and accuracy passes the 0.80 threshold. The +240/+180 MB at R=16 is the captured hidden_states buffer (replaces the captured logits + adds caching-allocator overhead of the eager tail). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The Full prefill backend's request-slot count was a module-level constant (FULL_CG_PREFILL_REQ_SLOTS = 16) with a TODO to promote it to a per-phase config knob. Operators tuning Full prefill CG could not change it without forking sglang. Promote it to PhaseConfig.req_slots, defaulting to 16 so existing behavior is unchanged. Tunable via the existing JSON config plumbing: --cuda-graph-config '{"prefill": {"backend": "full", "req_slots": 128}}' req_slots only matters for the Full backend; BCG (bs=1 only) and TC_PIECEWISE (bs-invariant via torch.compile) ignore it. The runner still caps min(req_slots, max_bs) so out-of-range values fall back to max_bs. No top-level engine arg — req_slots is a deep-tuning knob that only matters under --cuda-graph-backend-prefill=full, and users in that territory are already passing --cuda-graph-config. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Oasis-Git · 2026-06-16T05:04:35Z

/tag-run-ci-label

YazhiGao · 2026-06-16T07:37:38Z

why not directly use prefill_max_request for FULL_CG_PREFILL_REQ_SLOTS?

Oasis-Git · 2026-06-16T19:01:58Z

@YazhiGao It may introduce extra memory cost especially in flashinfer. Also users can set this hyperparameter through cuda config settings.

TcPiecewiseCudaGraphBackend's _run_compile_pass fires inside resolve_prefill_backend(self) and re-enters self.capture_prepare, which reads self._capture_req_slots. The attribute was previously assigned AFTER the resolve call, crashing on AttributeError at server startup for any TC_PIECEWISE-default config (e.g. test_basic_sanity with Llama-3.1-8B-Instruct). Pre-init to 1 (the non-Full default) before resolve_prefill_backend, then overwrite with the configured req_slots if the backend turns out to be Full. Same ordering trap the comment a few lines above already documented for self.layer_model. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Oasis-Git · 2026-06-18T21:12:00Z

GPU test all passed: https://github.com/sgl-project/sglang/actions/runs/27653872187/job/81801639686?pr=27988

…prefill # Conflicts: # python/sglang/srt/layers/attention/flashinfer_backend.py # python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py # python/sglang/srt/server_args.py

…ked prefill - Rename the prefill PhaseConfig knob req_slots -> full_prefill_max_req (specific name per review), updating ALLOWED_KEYS, the runner read, and docs. - Default it to None = auto: derive max(chunked_prefill_size // 512, 1) at consumption in PrefillCudaGraphRunner instead of a hardcoded 64. Verified on merged main: test_full_cuda_graph_prefill.py passes (Qwen3-8B, --cuda-graph-backend-prefill=full, mgsm_en 0.844 >= 0.80). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

ch-wan

LGTM

Oasis-Git requested review from Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, hebiao064, hnyls2002, ispobock and merrymercy as code owners June 12, 2026 02:43

Oasis-Git changed the title ~~[Experimental] Full Cuda Graph Support for Prefill~~ [Experimental][WIP] Full Cuda Graph Support for Prefill Jun 12, 2026

Oasis-Git force-pushed the cg-prefill branch 4 times, most recently from 2d0e47a to 5058a79 Compare June 16, 2026 04:26

Oasis-Git force-pushed the cg-prefill branch from 5058a79 to 23d0f91 Compare June 16, 2026 04:36

Oasis-Git and others added 2 commits June 16, 2026 04:39

Oasis-Git force-pushed the cg-prefill branch from 23d0f91 to 3d749d0 Compare June 16, 2026 04:45

Oasis-Git changed the title ~~[Experimental][WIP] Full Cuda Graph Support for Prefill~~ [Experimental] Full Cuda Graph Support for Prefill Jun 16, 2026

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Jun 16, 2026

Oasis-Git and others added 2 commits June 16, 2026 12:05

Merge branch 'main' into cg-prefill

18828fe

merrymercy reviewed Jun 20, 2026

View reviewed changes

Comment thread python/sglang/srt/model_executor/cuda_graph_config.py Outdated

Comment thread python/sglang/srt/model_executor/cuda_graph_config.py Outdated

Oasis-Git and others added 2 commits July 3, 2026 20:37

Merge branch 'main' of https://github.com/sgl-project/sglang into cg-…

e922b99

…prefill # Conflicts: # python/sglang/srt/layers/attention/flashinfer_backend.py # python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py # python/sglang/srt/server_args.py

Oasis-Git added run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fastfail labels Jul 3, 2026

Merge branch 'main' into cg-prefill

3aae892

ch-wan added the release-highlight Candidate PR for release note highlight label Jul 6, 2026

ch-wan approved these changes Jul 7, 2026

View reviewed changes

ch-wan merged commit 3cbb756 into sgl-project:main Jul 7, 2026
305 of 338 checks passed

Oasis-Git deleted the cg-prefill branch July 7, 2026 02:15

ch-wan added release-highlight Candidate PR for release note highlight and removed release-highlight Candidate PR for release note highlight labels Jul 9, 2026

amd-bot mentioned this pull request Jul 27, 2026

[CI Monitor] Daily Report - 2026-07-27 bingxche/sglang-ci-bot#148

Open

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[Experimental] Full Cuda Graph Support for Prefill (sgl-project#27988)

e39be58

This was referenced Sep 18, 2026

task(sglang): reconstruct Full Prefill CUDA Graph aitofound/ScienceIDE_Env#818

Closed

code(sglang): vendor full-prefill starting revision aitofound/ScienceIDE_Env#819

Closed

Navigation Menu

Repository navigation

[Experimental] Full Cuda Graph Support for Prefill - #27988

[Experimental] Full Cuda Graph Support for Prefill#27988
ch-wan merged 8 commits into
sgl-project:mainfrom
Oasis-Git:cg-prefill

Oasis-Git commented Jun 12, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jun 12, 2026

Uh oh!

Oasis-Git commented Jun 16, 2026

Uh oh!

YazhiGao commented Jun 16, 2026

Uh oh!

Oasis-Git commented Jun 16, 2026

Uh oh!

Oasis-Git commented Jun 18, 2026

Uh oh!

Uh oh!

Uh oh!

ch-wan left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

Oasis-Git commented Jun 12, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

gemini-code-assist Bot commented Jun 12, 2026

Uh oh!

Oasis-Git commented Jun 16, 2026

Uh oh!

YazhiGao commented Jun 16, 2026

Uh oh!

Oasis-Git commented Jun 16, 2026

Uh oh!

Oasis-Git commented Jun 18, 2026

Uh oh!

Uh oh!

Uh oh!

ch-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Oasis-Git commented Jun 12, 2026 •

edited by github-actions Bot

Loading