chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support releasing CUDA graph memory when paused (#7873)

Browse Source 
Co-authored-by: ryang-max <y1cunhui.yang@gmail.com>
Co-authored-by: ryang <38470282+ryang-max@users.noreply.github.com>
This commit is contained in:
fzyzcjy
2025-10-28 14:40:50 +08:00
committed by GitHub
parent d2b8c4123e
commit 691c8534cf
7 changed files with 94 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/constants.py
View File

@@ -1,3 +1,10 @@
# GPU Memory Types
GPU_MEMORY_TYPE_KV_CACHE = "kv_cache"
GPU_MEMORY_TYPE_WEIGHTS = "weights"
GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph"
GPU_MEMORY_ALL_TYPES = [
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
GPU_MEMORY_TYPE_CUDA_GRAPH,
]

4
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
View File

@@ -18,6 +18,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import
is_weak_contiguous,
)
from sglang.srt.distributed.parallel_state import in_the_same_node_as
from sglang.srt.environ import envs
from sglang.srt.utils import is_cuda, is_hip, log_info_on_rank0
logger = logging.getLogger(__name__)
@@ -210,6 +211,7 @@ class CustomAllreduce:
self.register_buffer(self.buffer)
self.disabled = False
self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
@staticmethod
def create_shared_buffer(
@@ -394,7 +396,7 @@ class CustomAllreduce:
if _is_hip:
return self.all_reduce_reg(input)
else:
return self.all_reduce(input, registered=True)
return self.all_reduce(input, registered=not self.tms_cudagraph)
else:
# If warm up, mimic the allocation pattern since custom
# allreduce is out-of-place.

3
python/sglang/srt/environ.py
View File

@@ -239,6 +239,9 @@ class Envs:
SGLANG_IMAGE_MAX_PIXELS = EnvInt(16384 * 28 * 28)
SGLANG_RESIZE_RESAMPLE = EnvStr("")
# Release & Resume Memory
SGLANG_MEMORY_SAVER_CUDA_GRAPH = EnvBool(False)
# Ktransformers
SGLANG_KT_MOE_NUM_GPU_EXPERTS = EnvInt(None)
SGLANG_KT_MOE_CPUINFER = EnvInt(None)

17
python/sglang/srt/managers/scheduler_update_weights_mixin.py
View File

@@ -5,7 +5,12 @@ from typing import TYPE_CHECKING, Tuple
import torch
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
from sglang.srt.constants import (
GPU_MEMORY_ALL_TYPES,
GPU_MEMORY_TYPE_CUDA_GRAPH,
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.srt.managers.io_struct import (
DestroyWeightsUpdateGroupReqInput,
DestroyWeightsUpdateGroupReqOutput,
@@ -104,7 +109,7 @@ class SchedulerUpdateWeightsMixin:
tags = recv_req.tags
if tags is None or len(tags) == 0:
tags = [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]
tags = GPU_MEMORY_ALL_TYPES
for tag in tags:
self.offload_tags.add(tag)
@@ -120,6 +125,9 @@ class SchedulerUpdateWeightsMixin:
torch.distributed.barrier(self.tp_cpu_group)
self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH)
return ReleaseMemoryOccupationReqOutput()
def resume_memory_occupation(
@@ -128,11 +136,14 @@ class SchedulerUpdateWeightsMixin:
tags = recv_req.tags
if tags is None or len(tags) == 0:
tags = [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]
tags = GPU_MEMORY_ALL_TYPES
for tag in tags:
self.offload_tags.remove(tag)
if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_CUDA_GRAPH)
if GPU_MEMORY_TYPE_WEIGHTS in tags:
self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
torch.distributed.barrier(self.tp_cpu_group)

14
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -21,12 +21,14 @@ import inspect
import logging
import os
from contextlib import contextmanager
from functools import partial
from typing import TYPE_CHECKING, Callable, Optional, Union
import torch
import tqdm
from torch.profiler import ProfilerActivity, profile
from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
from sglang.srt.custom_op import CustomOp
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -64,6 +66,7 @@ from sglang.srt.utils import (
require_mlp_tp_gather,
)
from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
try:
from kt_kernel import AMXMoEWrapper
@@ -518,7 +521,16 @@ class CudaGraphRunner:
logger.info(log_message)
def _capture_graph(self, graph, pool, stream, run_once_fn):
with self.device_module.graph(graph, pool=pool, stream=stream):
memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=self.model_runner.server_args.enable_memory_saver
and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
)
graph_fn = (
partial(memory_saver_adapter.cuda_graph, tag=GPU_MEMORY_TYPE_CUDA_GRAPH)
if memory_saver_adapter.enabled
else self.device_module.graph
)
with graph_fn(cuda_graph=graph, pool=pool, stream=stream):
out = run_once_fn()
return out

20
python/sglang/srt/utils/torch_memory_saver_adapter.py
View File

@@ -41,6 +41,12 @@ class TorchMemorySaverAdapter(ABC):
def region(self, tag: str, enable_cpu_backup: bool = False):
raise NotImplementedError
def cuda_graph(self, **kwargs):
raise NotImplementedError
def disable(self):
raise NotImplementedError
def pause(self, tag: str):
raise NotImplementedError
@@ -61,6 +67,12 @@ class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter):
def region(self, tag: str, enable_cpu_backup: bool = False):
return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup)
def cuda_graph(self, **kwargs):
return _memory_saver.cuda_graph(**kwargs)
def disable(self):
return _memory_saver.disable()
def pause(self, tag: str):
return _memory_saver.pause(tag=tag)
@@ -81,6 +93,14 @@ class _TorchMemorySaverAdapterNoop(TorchMemorySaverAdapter):
def region(self, tag: str, enable_cpu_backup: bool = False):
yield
@contextmanager
def cuda_graph(self, **kwargs):
yield
@contextmanager
def disable(self):
yield
def pause(self, tag: str):
pass