diff --git a/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py b/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py index e6fdfa8a7..073b031de 100644 --- a/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py +++ b/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py @@ -7,10 +7,10 @@ from torch.nn import functional as F from transformers import AutoConfig from sglang.benchmark.bench_utils import run_bench +from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( fused_moe as fused_moe_triton, ) -from sglang.srt.model_executor.cuda_graph_runner import set_torch_compile_config def get_model_config(model_name: str, tp_size: int): diff --git a/docs_new/docs/advanced_features/breakable_cuda_graph.mdx b/docs_new/docs/advanced_features/breakable_cuda_graph.mdx index 3497ba837..d41bbb5cd 100644 --- a/docs_new/docs/advanced_features/breakable_cuda_graph.mdx +++ b/docs_new/docs/advanced_features/breakable_cuda_graph.mdx @@ -33,7 +33,7 @@ This mode is intended for debugging only — it eliminates the performance benef For production use, you can mark specific functions as "non-graphable" using the `@eager_on_graph` decorator. During CUDA graph capture, these functions run eagerly between captured graph segments. Outside of capture, they behave normally. ```python -from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import eager_on_graph +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import eager_on_graph @eager_on_graph(enable=True) def my_dynamic_op(x): @@ -44,7 +44,7 @@ def my_dynamic_op(x): You can also insert a bare graph break (no computation) using the `break_graph()` helper: ```python -from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import break_graph +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import break_graph def forward(self, x): x = self.layer1(x) @@ -169,16 +169,16 @@ For typical use cases with a small number of graph breaks, the overhead is negli - python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py + python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py Core implementation: eager_on_graph, BreakableCUDAGraph, BreakableCUDAGraphCapture - python/sglang/srt/model_executor/breakable_cuda_graph/cuda_utils.py + python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py CUDA runtime binding utilities - python/sglang/srt/model_executor/cuda_graph_runner.py - Integration with the main CUDA graph runner + python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py + Integration with CUDA graph runner backends python/sglang/srt/server_args.py diff --git a/docs_new/docs/advanced_features/piecewise_cuda_graph.mdx b/docs_new/docs/advanced_features/piecewise_cuda_graph.mdx index 701bb9ae1..d1c3e7dd8 100644 --- a/docs_new/docs/advanced_features/piecewise_cuda_graph.mdx +++ b/docs_new/docs/advanced_features/piecewise_cuda_graph.mdx @@ -264,8 +264,8 @@ Use `--enforce-piecewise-cuda-graph` to skip all auto-disable checks (for testin - python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py - Main runner: init, capture, replay + python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py + Backend implementation: compile, capture, replay python/sglang/srt/compilation/compile.py diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py index 009e49139..3c27960bb 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py +++ b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py @@ -3,16 +3,16 @@ Public API (also reachable via the deeper module paths): - BreakableCUDAGraph, BreakableCUDAGraphCapture — capture/replay - eager_on_graph — decorator that marks a callable as a graph break + - break_graph — helper that inserts a bare graph break - enable_breakable_cuda_graph — context that flips the Breakable runtime flag - is_in_breakable_cuda_graph — runtime flag getter -The legacy model_executor/breakable_cuda_graph/ package is a -backwards-compat shim that re-exports from here. """ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import ( # noqa: F401 BreakableCUDAGraph, BreakableCUDAGraphCapture, + break_graph, eager_on_graph, ) from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( # noqa: F401