[PP] Support prefill CUDA graph proxy tensors

[YAMY1234]

Motivation

Pipeline-parallel prefill could not use the breakable prefill CUDA Graph path because PP stages expose globally indexed attention layers and exchange PPProxyTensors, while the graph runner previously assumed a non-PP layer list and rejected proxy inputs and outputs. This leaves genuine small aggregate forwards dominated by launch and control overhead. On Qwen3.5 PP4 with max_prefill_tokens == chunked_prefill_size == 2048, the eager path peaked at 34,857 tok/s at concurrency 16.

Modifications

• Keep attention and companion-layer lists aligned with global layer IDs and evaluate graph eligibility only on the local PP stage.
• Allocate and register stable graph buffers for PP hidden states, residuals, and optional top-k indices, then copy live proxy values into those buffers before replay.
• Pass proxy tensors through the captured stage forward and trim padded proxy outputs back to the real token count.
• Keep PP breakable prefill CUDA Graph explicit opt-in. When enabled, capture actual aggregate forwards up to 8K tokens by default and use eager above that threshold; explicit max_bs or bs settings still override the default.

This does not change scheduler admission, request grouping, chunk boundaries, or aggregate prefill budgets. A 2K configuration still executes at most 2K aggregate new tokens in one model forward.

Accuracy Tests

path	GSM8K score	correct
eager prefill	0.990	198/200
graph prefill	0.980	196/200

Both arms used the same 200 GSM8K five-shot chat prompts, 128 threads, max_tokens=16000, temperature=0.6, top_p=0.95, top_k=20, and runtime decode seed 42. The captured request payload files were byte-identical. Requests do not carry a sampling seed, so the observed 1.0 percentage-point difference is not treated as an accuracy change.

Targeted CPU tests cover global-layer alignment, stable PP proxy replay inputs (including hidden-only contracts), signature-safe proxy/embed binding, proxy output slicing, PP opt-in policy, and default and explicit capture limits.

Speed Tests and Profiling

AgentX end-to-end benchmark

The largest matched dual-axis end-to-end gain in the 900-second AgentX screen was at C44 on aggregate TP2 with B3/CG4. The source, model, container, Dynamo and AIPerf commits, MTP, cache policy, routing, chunking, and decode batch and graph caps were fixed. The effective runtime change was eager prefill versus breakable prefill CUDA Graph.

prefill path	job	Total TPS/GPU	P90 output TPS/user	P50 output TPS/user
eager	3411173	42,847.978	205.992	340.182
breakable CUDA Graph	3411823	48,733.914	245.690	374.802
change		+13.74%	+19.27%	+10.18%

The eager and graph arms retained 862 and 974 requests with zero request errors. The graph arm had 15 credits still in flight at the phase-end grace timeout; the eager arm had none. These are 900-second screening results, not 3,600-second submission numbers.

Prefill-only throughput benchmark

Configuration: GB300, Qwen3.5-397B-A17B NVFP4, one TP1×PP4 prefill worker and one TP4 decode worker, 8K ISL, 1 OSL, radix cache off, MTP off, and max_prefill_tokens == chunked_prefill_size. Throughput is aggregate input-plus-output tokens per second; all points completed 10 × concurrency requests with exact 8192/1 token lengths and zero request errors.

Each concurrency cell is eager → graph tok/s. Both paths use the same source and differ only in cuda-graph-backend-prefill=disabled/breakable.

aggregate cap	C16	C32	C64	C128	C256	eager peak	graph peak	peak change
2K	34,857 → 86,476	34,852 → 86,561	33,274 → 86,433	33,942 → 86,100	33,660 → 85,912	34,857 @ C16	86,561 @ C32	2.48× / +148.33%
4K	65,670 → 127,903	66,194 → 126,598	65,977 → 127,942	64,931 → 127,732	64,909 → 127,804	66,194 @ C32	127,942 @ C64	1.93× / +93.28%
8K	131,807 → 146,105	131,714 → 149,428	134,906 → 150,121	131,560 → 151,261	131,703 → 151,349	134,906 @ C64	151,349 @ C256	1.12× / +12.19%
16K	129,637 → 115,657	174,196 → 169,433	177,064 → 171,809	178,023 → 173,620	177,080 → 174,068	178,023 @ C128	174,068 @ C256	0.98× / -2.22%
32K	94,412 → 74,474	143,688 → 124,359	186,020 → 177,690	186,828 → 180,068	186,224 → 181,276	186,828 @ C128	181,276 @ C256	0.97× / -2.97%

The benefit is concentrated in small forwards, so PP uses breakable prefill CUDA Graph only when the backend is selected explicitly. Its default capture limit is 8K actual aggregate tokens, the last measured size with a positive peak result; larger forwards automatically use eager. Operators can override the limit with --cuda-graph-max-bs-prefill or explicit capture buckets. Fixed-2K peak rises from 18.66% of eager 32K peak to 47.75% of graph 32K peak, a 2.56× recovery of the relative small-batch throughput.

At matched C64, fixed 2K improves from 33,274 to 86,433 tok/s, or 2.60×. The graph sweep logged 95,596 full 2K forwards as graph replays and zero as eager; no forward exceeded the 2K aggregate cap.

Stage profiling

PP rank	eager run_batch median	graph run_batch median	eager kernel busy/step	graph kernel busy/step
0	73.06 ms	19.82 ms	18.56 ms	18.34 ms
1	75.24 ms	21.29 ms	17.40 ms	17.19 ms
2	73.86 ms	20.10 ms	16.29 ms	16.20 ms
3	56.12 ms	16.27 ms	12.23 ms	12.15 ms

The graph reduces stage wall time by 71–73% while kernel busy time stays within 1.2%, consistent with removing launch and control bubbles without changing model compute. In the common trace window, average kernel-active GPUs increased from 0.709/4 to 2.009/4. Profiler overhead affects cadence, so throughput numbers come only from unprofiled runs.

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

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. 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
4. 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 #33265642901
Latest PR Test (Extra): ❌ Run #33265642817
Latest PR Test (AMD ROCm 7.2): ❌ Run #33265642882

[YAMY1234]

/tag-and-rerun-ci

[YAMY1234]

Three follow-up correctness and memory issues to address before merge.

[YAMY1234]

/rerun-failed-ci bypass-fastfail a0f92e8

[Oasis-Git]

CI Pass: https://github.com/sgl-project/sglang/actions/runs/33265642901/job/99135188885?pr=36248