Skip to content

[Perf] cuda-graph: Raise MLA prefill capture default to chunked size capped at 8192 #45565

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Fridge003 merged 4 commits into from
Oct 22, 2026
Merged

[Perf] cuda-graph: Raise MLA prefill capture default to chunked size capped at 8192 #45565

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d497ebf
Remove the MLA-specific prefill CUDA graph capture cap
Fridge003 Oct 21, 2026
4a6a632
Merge branch 'main' into codex/revert-mla-prefill-capture-limit
Fridge003 Oct 21, 2026
b34e950
Cap the default prefill CUDA graph capture size at 8192
Fridge003 Oct 21, 2026
264cfe3
Preserve non-MLA prefill capture defaults
Fridge003 Oct 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions python/sglang/srt/arg_groups/memory_hook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,13 +231,11 @@ def handle_gpu_memory_settings(server_args: Any):
decode_cuda_graph_config.max_bs = cfg.torch_compile_max_bs

if prefill_cuda_graph_config.max_bs is None:
# Refer to pr #15927, by default we set the prefill max_bs to the chunked prefill size.
# For MLA backend, the introduction of piecewise cuda graph will influence the kernel dispatch difference compared to the original mode.
# To avoid the performance regression, we set max_bs to 2048 by default.
# Default to the chunked prefill size; cap MLA capture at 8192 tokens.
if not use_mla_backend(server_args):
prefill_cuda_graph_config.max_bs = cfg.chunked_prefill_size
else:
prefill_cuda_graph_config.max_bs = 2048
prefill_cuda_graph_config.max_bs = min(8192, cfg.chunked_prefill_size)

# For opt-in PP breakable graphs, capture small aggregate-token
# buckets by default and leave larger forwards on the eager path.
Expand Down
Loading