Repository navigation
Size KV pool after CUDA graph capture (opt-in) - #30157
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a post-capture KV cache sizing feature using CUDA Virtual Memory Management (VMM) to dynamically size and physically back the KV cache pool based on measured free memory after CUDA graph capture. While the implementation is a significant optimization, several critical issues must be addressed: a missing error check on cuMemAddressReserve, a multi-process race condition during inline compilation due to non-unique library names, a unit mismatch when querying device memory capacity (bytes vs. megabytes) that wastes memory, an incorrect slot calculation for HND (Paged) KV cache leading to severe over-allocation, and an unsafe deallocation order when clearing buffers.
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.
|
/tag-and-rerun-ci |
f9782e6 to
fda9465
Compare
Add SGLANG_ENABLE_POST_CAPTURE_KV_SIZING (default off). When enabled, the KV pool is reserved as CUDA VMM virtual memory, CUDA graphs are captured against its stable pointers, then the pool is sized from measured free memory and physically backed via cuMemMap -- dropping the heuristic per-token activation reserve baked into mem_fraction_static. When the flag is off, behavior is unchanged: the feature is fully gated behind ServerArgs.post_capture_kv_sizing_planned(), and the mem_fraction_static heuristic (including reserve_for_graph_mb / reserve_for_deepep_a2a_mb) is byte-for-byte equivalent to the previous logic. Adds test/registered/mem_cache/test_post_capture_kv_sizing.py guarding the feature (server comes up, sizing path runs, pool sized, gsm8k unchanged).
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
ace5952 to
c1ad7d7
Compare
Post-capture sizing measured only KV-pool-free memory and left near-zero headroom, so the pool over-sized and inference OOM'd on transient activation (eager decode above the captured graph max_bs, and prefill/logits allocations that the captured graphs do not hold resident). Floor the resize headroom to the activation working set via pre_capture_activation_reserve_mb() (same reserve used pre-capture), and apply it for all models, not just mamba. Renamed from mamba_pre_capture_reserve_mb since it covers the generic activation reserve.
At post-capture resize, keep the minimal mem_fraction headroom by default and floor it to mamba_pre_capture_reserve_mb only when decode capture max_bs < max_running_requests (eager decode above the captured graph) or a mamba state cache is present; warn on the eager-decode case. Restore the per-arena stub suffix to the instance counter so all TP ranks emit identical stub sources and the shared build directory stays idempotent; the pid-based suffix made each rank's source unique and raced on the shared build dir.
7ed5661 to
4a96aa6
Compare
Resolve server_args.py conflict: keep the PR's post-capture KV sizing refactor (reserve_for_graph_mb / reserve_for_deepep_a2a_mb helpers) and port main's resolved-view fixes (_resolved().enable_dp_attention and resolved_view(self).moe_a2a_backend) into those helpers.
|
/rerun-test test/registered/mem_cache/test_post_capture_kv_sizing.py |
|
Results for 🚀 |
| if eager_decode_gap or self.mambaish_config is not None: | ||
| headroom_gb = max( | ||
| headroom_gb, | ||
| self.server_args.mamba_pre_capture_reserve_mb( |
There was a problem hiding this comment.
this fallback is too relaxed.
- For a decode engine, it reserves > 10 GB, which is too much
- why do we call "mamba" here?
There was a problem hiding this comment.
- this 10GB reservation happens only when
max_running_reqsis not set and decode can run at a batch size over the captured size. this should not be a common case - good catch, we need to change this name to be generic
| # Derive from the real buffers when they exist (covers arbitrary layouts, | ||
| # e.g. vectorized_5d); fall back to _kv_buffer_shapes for the pre-allocation | ||
| # post-capture call, which only runs for NHD/HND. | ||
| if getattr(self, "k_buffer", None) and getattr(self, "v_buffer", None): |
There was a problem hiding this comment.
do not use getattr anymore!
There was a problem hiding this comment.
i will refactor this to add the attribute always
|
|
||
| @property | ||
| def post_capture_active(self) -> bool: | ||
| return getattr(self.full_kv_pool, "post_capture_active", False) |
There was a problem hiding this comment.
do not use getattr anymore!
There was a problem hiding this comment.
i will remove this one
2 participants
Problem: the KV pool is sized before CUDA-graph capture, so graph memory is a heuristic guess baked into
mem_fraction_static. On large configs the guess can be off by ~10 GB/GPU; the padded activation reserve silently absorbs the difference, wasting memory on some configs and risking OOM on others.Fix: add
SGLANG_ENABLE_POST_CAPTURE_KV_SIZING(opt-in, default off). When enabled:cuMemMap), and drop the per-token activation reserve.Test
test/registered/mem_cache/test_post_capture_kv_sizing.pylaunches a server with the flag on and asserts:/server_inforeportsmax_total_num_tokens > 0,Registered in the base CI tier (
base-b/1-gpu-large) so it runs on every PR.CI States
Latest PR Test (Base): ❌ Run #28883219805
Latest PR Test (Extra): ❌ Run #28883219626