Share BCG output buffers across capture sizes - #27659
Merged
Merged
Conversation
Reuse one maximum-size output buffer while capturing breakable CUDA graph token buckets, and store sliced views for smaller capture sizes. Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
merrymercy
requested review from
Fridge003,
Ying1123,
hnyls2002 and
ispobock
as code owners
June 9, 2026 08:16
Contributor
Author
|
/tag-and-rerun-ci |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a shared output buffer mechanism in breakable_cuda_graph_runner.py to optimize CUDA graph capture across different token sizes, adding helper methods _slice_output and _copy_output_to_buffer to manage output slicing and copying. The review feedback correctly identifies potential TypeError issues if the model's forward pass returns None for optional outputs, and suggests robustly handling None values in both helper methods.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Contributor
Author
|
/tag-and-rerun-ci |
Fridge003
approved these changes
Jun 10, 2026
Oasis-Git
added a commit
to Oasis-Git/sglang
that referenced
this pull request
Jun 11, 2026
…class
The cg-refactor's cuda-graph dict key was `int | str` produced by
_make_graph_key(bs, stream_idx, variant_label) — e.g. `bs` for plain
decode, `f"{stream_idx}_{bs}"` for pdmux, `f"{lora}_{stream_idx}_{bs}"`
for LoRA-variant capture. The str form was opaque and conflated three
independent axes.
Replace with a frozen dataclass:
@DataClass(frozen=True)
class ShapeKey:
size: int # bs * num_tokens_per_bs (decode)
# or num_tokens (prefill)
stream_idx: Optional[int] = None
variant_label: Optional[str] = None
`size` fuses bs and num_tokens_per_bs (matches the existing convention
that bs and token count are unified). `frozen=True` gives us __eq__ /
__hash__ for free, so it drops into _graphs / _outputs dict lookups
unchanged.
This also lets BCG's output-buffer sharing (sgl-project#27659) read shape_key.size
directly — the redundant num_tokens kwarg added to capture_one in the
previous commit goes away, and the abstract surface across all
backends (BCG / Full / TcPiecewise / NPU) is back to a single typed
shape_key argument.
Touches:
- new python/sglang/srt/model_executor/runner/shape_key.py
- decode / prefill runners construct ShapeKey at capture and replay
- 3 EAGLE speculative runners override _make_graph_key to return ShapeKey
- 4 backends (base, BCG, Full, TcPiecewise, NPU) type shape_key: ShapeKey
- BCG capture_one drops num_tokens kwarg; reads shape_key.size
Verified on gpt-oss-20b, 1x B200, --chunked-prefill-size 16384,
--cuda-graph-backend-prefill breakable: prefill BCG capture cost still
2.00 GB (matches the 1.99 GB from the pre-refactor bcg-slice commit
within noise; baseline on main is 2.68 GB). Server serves /generate.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Oasis-Git
added a commit
to Oasis-Git/sglang
that referenced
this pull request
Jun 11, 2026
The previous ShapeKey refactor commit changed the decode + EAGLE _make_graph_key methods to construct ShapeKey(size=bs * num_tokens_per_bs). That subtly changed the dict-key identity vs the original _make_graph_key(bs, stream_idx, variant_label) which keyed on `bs` alone. Revert to ShapeKey(size=bs) so the cuda-graph key identity is identical to pre-refactor behavior. BCG output-buffer sharing (sgl-project#27659) is unaffected: prefill constructs ShapeKey(size=num_tokens) which IS the leading dim of the captured output, and decode BCG is not the default code path. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Oasis-Git
added a commit
to Oasis-Git/sglang
that referenced
this pull request
Jun 11, 2026
Strip the # PR sgl-project#27659 ... explanatory blocks from BreakableCudaGraph- Backend (capture iteration order, copy-into-shared-buffer mechanics) and the tombstone left where the old _make_graph_key helper used to live. Simplify shape_key.py's module docstring to its one-line subject. The code is self-explanatory; the historical context belongs in the commit / PR description, not in source where it rots. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
jakki-amd
pushed a commit
to jakki-amd/sglang
that referenced
this pull request
Sep 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Noneoutputs when sharing captured output buffers.Verification
git diff --check -- python/sglang/srt/model_executor/breakable_cuda_graph_runner.pypython3 -m py_compile python/sglang/srt/model_executor/breakable_cuda_graph_runner.pyblack --check python/sglang/srt/model_executor/breakable_cuda_graph_runner.pyruff check --select=F401,F821 python/sglang/srt/model_executor/breakable_cuda_graph_runner.pyOriginal commits
a2c14728ae0747b8eeCI States
Latest PR Test (Base): ✅ Run #27193490863
Latest PR Test (Extra): ❌ Run #27193490917