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

[nvpohanh]

[by Codex]

Motivation

When --context-length is smaller than a prefill CUDA-graph capture bucket, the old capture path creates one synthetic request with seq_len=num_tokens. That exceeds the per-request context limit and can cause the #21112 illegal-memory-access failure during graph warm-up.

Clamping the aggregate prefill token ceiling to context_length avoids that crash, but unnecessarily disables graph replay for legal multi-request prefills whose total token count exceeds one request's context length.

Design

This PR fixes the capture shape rather than applying a per-request cap to an aggregate-token setting:

• Resolve legal prefill capture buckets once in ModelRunner.init_prefill_cuda_graph() and write the filtered list back to cuda_graph_config.prefill.bs before constructing the runner.
• Keep every backend's capture buckets expressed in aggregate tokens and cap them at context_length * max_capture_requests: request-pool size for tc_piecewise/breakable, or the resolved fixed request-slot count for full.
• During capture, partition a bucket of num_tokens into the fewest synthetic requests such that every request has seq_len <= context_length.
• For full, preserve its fixed request-axis contract by padding any unused request slots with zero-length sentinels after the context-bounded synthetic requests.
• If filtering leaves no legal bucket, log the backend, context length, and request-pool capacity and fall back to eager prefill.
• Remove the earlier ServerArgs-time context_length * max_running_requests calculation because max_running_requests is not always resolved at that stage.
• For DeepSeek's MHA prefill path under BCG, register and select the existing attn_mha companion on CUDA as well as HIP. This ensures the captured attention op receives MHA head metadata instead of the same-layer-id attn_mqa metadata.

This preserves graph coverage for multi-request prefills while keeping every synthetic request within the per-request context bound.

Validation

• git diff --check
• python -m py_compile on all five changed Python files
• Black check in lmsysorg/sglang:dev-cu13: all five changed files left unchanged

4xGB200 DSR1 NVFP4 verification

Passed on four NVIDIA GB200 GPUs at commit aec6c63161 using a staged DeepSeek-R1-0528 ModelOpt NVFP4 checkpoint:

• TP4, FlashInfer attention, context_length=2048, max_running_requests=8, chunked_prefill_size=4096, decode CUDA graphs disabled, and FlashInfer autotuning disabled.
• Prefill graph configuration: backend=breakable, bs=[4096]. The only capture bucket is therefore larger than one request's context and is constructed as two legal synthetic requests.
• All four ranks completed Capture target prefill CUDA graph for 4,096 tokens in about 20.9 seconds; capture used about 1.41 GiB per GPU.
• The server reached readiness. A two-prompt /generate request returned four completion tokens for each prompt; server-reported prompt lengths were 1,301 and 1,301 tokens (2,602 aggregate, greater than the 2,048-token per-request context).
• Both runtime prefills logged cuda graph: True.
• AIHub Slurm step 4272943.2 completed with exit code 0:0.

Earlier H100 verification

Passed on one H100 80GB at commit e689f496ba:

• Qwen/Qwen2.5-1.5B-Instruct with --context-length 2048 --mem-fraction-static 0.7 --attention-backend flashinfer --cuda-graph-backend-prefill tc_piecewise.
• Piecewise prefill CUDA-graph capture completed across every configured bucket through 8,192 aggregate tokens (4 * context_length).
• The server reached readiness and completed its startup /generate request successfully.

Related

• Fixes the capture-shape issue underlying [Bug] Piecewise CUDA Graph crashes with illegal memory access on H100 (FA3 backend) during warmup_compile #21112 / fix(server): clamp piecewise_cuda_graph_max_tokens to context_length #22516.
• Follow-up to the discussion in fix(server): capture legal multi-request prefill CUDA graph batches #30206.

---

CI States

Latest PR Test (Base): ✅ Run #30425272323
Latest PR Test (Extra): ❌ Run #30425272188

[gemini-code-assist]

Warning

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

[nvpohanh]

[by Codex]
/tag-and-rerun-ci

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

cc @janbernloehr @hnyls2002 for vis

[janbernloehr]

Note: I'm currently traveling without my laptop, so I asked Claude to dig into this — the analysis below was created with Claude Code, and I haven't been able to run anything locally myself.

Thanks for the PR — the perf observation is real (with the current clamp, a multi-request prefill batch whose total tokens exceed context_length falls back to eager even though every individual request is legal). Unfortunately the fix as written doesn't work, for two reasons:

1. TypeError at startup in the default config

max_running_requests defaults to None in ServerArgs and is never assigned a computed value in server_args.py — the effective value is derived later in the scheduler from KV-cache capacity (the only other use in this file guards it with or, see the disaggregation_mode == "decode" branch below). So

self.context_length * self.max_running_requests

raises TypeError: unsupported operand type(s) for *: 'int' and 'NoneType' for anyone passing --context-length without --max-running-requests.

2. Even with max_running_requests set, this reintroduces the #21112 crash

The clamp exists because of how prefill graph capture builds its dummy batch, not because of what replay batches can contain. PrefillCudaGraphRunner.capture_prepare puts all captured tokens into a single fake request:

bs = 1
"seq_lens": torch.tensor([num_tokens], device=self.device),

so the attention backend's metadata init sees one sequence of length num_tokens. The buffers that overflowed in #21112 are sized per request by context_length (the req_to_token row that FA3 builds its page table from). The aggregate capacity context_length * max_running_requests only exists across rows of the req pool — a single-request capture batch can never legally use it.

Concretely: on H100 with --context-length 2048 and a non-MLA model, prefill.max_bs defaults to chunked_prefill_size = 8192. The product cap (2048 × anything ≥ 4) never binds, so warmup again compiles a single 8192-token dummy request against 2048-wide buffers — the exact #21112 configuration. In practice the multiplied clamp is a no-op, i.e. equivalent to reverting #22516.

(Minor additional wrinkle: for the breakable prefill backend, max_bs is a request count, so clamping it by a token product mixes units — the existing clamp already has this issue, but the multiplication compounds it.)

Suggested direction

The bottleneck is the capture path, so the fix that would actually unlock this is in capture_prepare: split the dummy tokens across multiple fake requests when num_tokens > context_length — bs = ceil(num_tokens / context_length), each seq_len ≤ context_length, with bs bounded by the req-pool size and the per-request static buffers sized accordingly. Once capture can represent that shape legally, the clamp can be relaxed to context_length × <number of dummy requests capture can build>. For tc_piecewise this is cheap to validate since attention runs eagerly outside the graph — capture just has to execute once without OOB.

Until then I think the conservative clamp has to stay, since the alternative is trading an eager-fallback slowdown for an illegal-memory-access crash at startup.

[nvpohanh]

I think I know what might be the issue.

First of all, prefill_cuda_graph_config.max_bs despite the name has "bs", it does NOT refer to the number of requests. Instead, it refers to the number of tokens in prefill. A better name would have been max_num_tokens.

Setting that aside, the main issue here is: what kind of dummy input batch does SGLang use to capture the PCG/BCG for prefill? Currently, it appears that SGLang just sends in one request with seqLen=num_tokens. But that is wrong if context-length is set to be smaller than max_num_tokens because context-length says that the max seqLen for each request is at most context-length, so the graph capture will run into issues.

The original PR #22516 fixed this issue by lowering the max_num_tokens (aka max_bs) setting to match context-length, but that is too restrictive when prefill can run with multiple requests.

The "correct" fix seems to be relaxing the max_num_tokens to context-length * max_running_requests, but when SGLang sends in dummy input batch for BCG/PCG graph capturing, it should be aware of context length and should NOT use a request with seqLen>context_lengh. Instead, it should send in multiple requests each each has SeqLen<=context_length for BCG/PCG graph capturing. The total number of tokens still match max_num_tokens (max_bs) but each request respects the context-length setting.

[nvpohanh]

[by Codex]
Updated to address both review findings:

• removed the ServerArgs-time context_length * max_running_requests calculation, so an unset request limit cannot raise TypeError;
• capture now partitions each token bucket into legal dummy requests with seq_len <= context_length, bounded by the request-pool capacity.

/tag-and-rerun-ci

[nvpohanh]

[by Codex]
Applied the Black formatting reported by lint and pushed the update.

/tag-and-rerun-ci

[janbernloehr]

(Still traveling without my laptop — this follow-up review was again put together with Claude Code.)

Thanks for the rework — this is the right approach now, and the tc_piecewise path looks correct to me: the static per-request buffers are allocated at req_to_token_pool.size, so the multi-request slices fit; the bucket filter guarantees the bs <= self.max_bs assert; and the start-loc/seq-len bookkeeping is consistent. I also like that it keys off model_config.context_len rather than only a user-set --context-length — that covers natively-small-context models too, which the hardcoded Llama-2 4096 clamp in server_args.py was working around (that special case could probably be removed in a follow-up).

Two remaining points before I'd consider this mergeable:

1. The "no bucket fits" path warns, then crashes anyway. init_prefill_cuda_graph (model_runner.py) only skips runner construction when the configured cuda_graph_config.prefill.bs is empty. With this PR, the configured list can be non-empty while the new runner-side filter empties self.capture_num_tokens. Construction then continues, and TcPiecewiseCudaGraphBackend._run_compile_pass unconditionally calls

cuda_graph_runner._run_dummy_forward(
    num_tokens=cuda_graph_runner.capture_num_tokens[0]
)

which raises IndexError on the empty list at startup — right after the new "Disable prefill CUDA graph capture…" warning. The max_num_tokens = 0 buffer registry is also an untested path. Could this fall back to the eager runner cleanly instead — e.g., run the filter before constructing the runner (in init_prefill_cuda_graph), or return early to eager_runner when the filtered list is empty?

2. The Breakable backend's bs=1 capture contract. The multi-request split applies to both prefill backends, but the BCG replay path explicitly documents "the captured graph is a bs=1 replay of layer_model.forward" (see execute() and the can_run_graph comment), and DSV4's opt-in path stashes per-bucket attention metadata built from the capture-time dummy batch. With a bucket larger than context_length, captured BCG segments would now bake static-buffer views of shape [k] instead of [1], and the stashed metadata comes from a k-request dummy batch. Has this been exercised with the Breakable backend and context_length smaller than a bucket? If not, I'd suggest restricting the multi-request split to tc_piecewise and keeping BCG buckets capped at context_length for now — that still fully fixes the #21112 case this PR targets.

[nvpohanh]

[by Codex] Addressed in 8b1b4ee.

1. init_prefill_cuda_graph now applies the same backend-aware bucket filter before constructing PrefillCudaGraphRunner. If no bucket survives, it logs the warning and routes prefill to eager_runner, so TcPiecewiseCudaGraphBackend never indexes an empty capture list.

2. Aggregate multi-request dummy capture is now limited to tc_piecewise. Breakable buckets remain capped at context_length, and its capture preparation remains bs=1, preserving the replay and capture-time metadata contract.

Validated with git diff --check and Python syntax compilation for both changed modules.

[nvpohanh]

   # Breakable captures a bs=1
  # graph and must retain the per-request context-length cap.

I don't like this. Could we make breakable also captures with multiple requests? If not, why can't we do so?

[nvpohanh]

[by Codex] Addressed in 6d07765. Breakable prefill now captures the same synthetic multi-request batches as tc_piecewise, and both backends admit buckets through request_pool_size × context_length.

This is safe because Breakable already owns stable request-axis buffers; attention runs at Breakable graph boundaries with replay-time metadata, and the request-layout-dependent output tail remains eager. The captured transformer segments are token-axis operations, so replay continues to support the live request layout.

Validated with git diff --check and Python syntax compilation for both changed modules.

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

All NV pipelines have passed. @Oasis-Git could you help us to review this since this is related to BCG? thanks!

[Oasis-Git]

LGTM.

[Oasis-Git]

All CI Passed: https://github.com/sgl-project/sglang/actions/runs/30425272323/job/90556547107?pr=30206