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

Piecewise Cuda Graph Support for gpt-oss model (#13045)

Browse Source 
Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>
This commit is contained in:
Yuwei An
2025-11-14 17:28:00 -08:00
committed by GitHub
parent 14dc052382
commit eae59b337e
5 changed files with 28 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
View File

@@ -26,7 +26,10 @@ import tqdm
from sglang.srt.compilation.compilation_config import CompilationConfig
from sglang.srt.compilation.compile import install_torch_compiled, set_compiled
from sglang.srt.compilation.piecewise_context_manager import set_forward_context
from sglang.srt.compilation.piecewise_context_manager import (
enable_piecewise_cuda_graph,
set_forward_context,
)
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 (
@@ -55,22 +58,6 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
_in_piecewise_cuda_graph = False
def is_in_piecewise_cuda_graph():
return _in_piecewise_cuda_graph
@contextmanager
def enable_piecewise_cuda_graph():
global _in_piecewise_cuda_graph
_in_piecewise_cuda_graph = True
yield
_in_piecewise_cuda_graph = False
@contextmanager
def freeze_gc(enable_cudagraph_gc: bool):