[Utils] Refactor device cache emptying#24861
Merged
Merged
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
ByronHsu
approved these changes
May 9, 2026
Collaborator
|
/tag-and-rerun-ci |
Collaborator
|
/tag-and-rerun-ci |
ByronHsu
pushed a commit
to ByronHsu/sglang
that referenced
this pull request
May 10, 2026
Replace direct torch.cuda.empty_cache() / memory_reserved() calls in continue_generation with the empty_device_cache() helper from sgl-project#24861, making the in-place pause resume path work on all device backends. Co-authored-by: Cursor <cursoragent@cursor.com>
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
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.
Motivation
SGLang has several paths that empty the PyTorch device allocator cache while also using
flush_cacheto clear internal memory pools such as KV cache and Mamba cache. Some scheduler paths still hard-codetorch.cuda.empty_cache(), which makes the allocator-emptying behavior CUDA-specific even though SGLang supports other device backends.This PR keeps the existing external API behavior while making the internal distinction clearer:
flush_cacheclears SGLang memory pools.empty_device_cacheonly releases unused cached blocks from the active device allocator.Modifications
empty_device_cache()as a small common helper for backend allocator cache emptying.flush_cacheand idle periodic cache emptying instead of directly callingtorch.cuda.empty_cache().get_available_gpu_memoryfor CUDA, XPU, NPU, and MUSA empty-cache paths.flush_cache_after_weight_update.flush_cachedocstring around memory pools such as KV cache and Mamba cache.Accuracy Tests
Not applicable. This PR does not change model forward behavior or numerical outputs.
Speed Tests and Profiling
Not applicable. This is a small cache-management refactor and preserves existing defaults.
Validation
python3 -m py_compile python/sglang/srt/utils/common.py python/sglang/srt/managers/scheduler.py python/sglang/srt/managers/scheduler_update_weights_mixin.py python/sglang/srt/managers/io_struct.pygit diff --checkNote: local pytest with
uv --directory python --extra testcould not run on macOS arm64 becausesgl-deep-gemm==0.0.1has no wheel for this platform.