Skip to content

GitLab

Unverified Commit ff0b64e1 authored by fzyzcjy's avatar fzyzcjy Committed by GitHub
Browse files

Ensure GPU work is finished when release memory occupation call is finished (#12592)

parent d84790db
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
python/sglang/srt/managers/scheduler_update_weights_mixin.py
View file @ ff0b64e1
...@@ -132,6 +132,8 @@ class SchedulerUpdateWeightsMixin: ...@@ -132,6 +132,8 @@ class SchedulerUpdateWeightsMixin:
if GPU_MEMORY_TYPE_CUDA_GRAPH in tags: if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH) self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH)
torch.cuda.synchronize()
return ReleaseMemoryOccupationReqOutput() return ReleaseMemoryOccupationReqOutput()
def resume_memory_occupation( def resume_memory_occupation(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment