diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index 6c3dcf67a..c64d0086a 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -76,8 +76,6 @@ def handle_pd_disaggregation(server_args: ServerArgs) -> None: server_args.disaggregation_transfer_backend != "fake" ), "Prefill server does not support 'fake' as the transfer backend" - server_args.disable_cuda_graph = True - if server_args.disaggregation_mode in ("prefill", "decode"): if ( envs.SGLANG_DISAGG_STAGING_BUFFER.get() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6a874517b..31af39d47 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3335,11 +3335,10 @@ class ServerArgs: # ------------------------------------------------------------------ # CUDA graph configuration resolution # ------------------------------------------------------------------ - # TODO: add unit tests in test/srt/test_server_args.py covering the - # precedence cascade + auto-disable matrix (follow-up PR). def _handle_cuda_graph_config(self): self._parse_cuda_graph_config() self._apply_cuda_graph_compatibility() + self._apply_cuda_graph_disaggregation_roles() self._validate_cuda_graph_config() # Warn on the final resolved config (not inside the compat cascade — # that path is skipped when the user explicitly sets the backend, @@ -3429,6 +3428,14 @@ class ServerArgs: elif self.cuda_graph_config.prefill.backend == Backend.FULL: self._disable_full_prefill_cudagraph_if_incompatible() + def _apply_cuda_graph_disaggregation_roles(self): + if self.disaggregation_mode == "prefill": + if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked: + self.cuda_graph_config.decode.backend = Backend.DISABLED + elif self.disaggregation_mode == "decode": + if (Phase.PREFILL, "backend") not in self._cuda_graph_config_locked: + self.cuda_graph_config.prefill.backend = Backend.DISABLED + def _disable_tc_piecewise_cudagraph_if_incompatible(self): from sglang.srt.arg_groups.overrides import resolved_view as _resolved_view diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 29c540579..28c05ecb7 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -8,6 +8,7 @@ from sglang.srt.distributed import get_tp_group from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -287,7 +288,9 @@ class DFlashWorkerV2(BaseSpecWorker): self._draft_worker.init_attention_backends() def init_cuda_graphs(self): - capture_decode_cuda_graph = not self.server_args.disable_cuda_graph + capture_decode_cuda_graph = ( + self.server_args.cuda_graph_config.decode.backend != Backend.DISABLED + ) if is_cuda() and capture_decode_cuda_graph: available_mem = get_available_gpu_memory(self.device, self.gpu_id) if available_mem < 1.0: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index acd6047f8..591260089 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1154,7 +1154,7 @@ class EAGLEWorkerV2(BaseSpecWorker): self.adaptive_controller.init_states( cuda_graph_bs=( None - if self.server_args.disable_cuda_graph + if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED) else self.server_args.cuda_graph_bs_decode ), ) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 74605e66c..f365c3867 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -13,6 +13,7 @@ from sglang.srt.layers.cp.base import is_cp_enabled, is_interleave from sglang.srt.model_executor.cuda_graph_config import ( Backend, CudaGraphConfig, + Phase, PhaseConfig, ) from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args @@ -1269,6 +1270,60 @@ class TestCudaGraphConfigDataclassAccess(CustomTestCase): self.assertEqual(config.compiler, "eager") +class TestCudaGraphDisaggregationRoles(CustomTestCase): + def _handled_args(self, **overrides): + args = ServerArgs(model_path="dummy", **overrides) + args.model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]), + is_piecewise_cuda_graph_disabled_model=False, + is_multimodal=False, + is_multimodal_piecewise_cuda_graph_supported=False, + ) + with ( + patch("sglang.srt.utils.is_cuda", return_value=True), + patch.object(ServerArgs, "use_mla_backend", return_value=False), + ): + args._handle_cuda_graph_config() + return args + + def test_cuda_graph_prefill_role_defaults_disable_decode_graph(self): + args = self._handled_args(disaggregation_mode="prefill") + + self.assertFalse(args.disable_cuda_graph) + self.assertEqual(args.cuda_graph_config.decode.backend, Backend.DISABLED) + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE) + + def test_cuda_graph_decode_role_defaults_disable_prefill_graph(self): + args = self._handled_args(disaggregation_mode="decode") + + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED) + self.assertNotEqual(args.cuda_graph_config.decode.backend, Backend.DISABLED) + + def test_cuda_graph_global_disable_still_disables_both_phases_for_all_roles(self): + for disaggregation_mode in ("prefill", "decode", "null"): + with self.subTest(disaggregation_mode=disaggregation_mode): + args = self._handled_args( + disaggregation_mode=disaggregation_mode, + disable_cuda_graph=True, + ) + + self.assertEqual( + args.cuda_graph_config.decode.backend, Backend.DISABLED + ) + self.assertEqual( + args.cuda_graph_config.prefill.backend, Backend.DISABLED + ) + + def test_cuda_graph_explicit_decode_backend_survives_prefill_role(self): + args = self._handled_args( + disaggregation_mode="prefill", + cuda_graph_backend_decode=Backend.FULL, + ) + + self.assertEqual(args.cuda_graph_config.decode.backend, Backend.FULL) + self.assertIn((Phase.DECODE, "backend"), args._cuda_graph_config_locked) + + class TestCutedslMoeMaxNumTokens(CustomTestCase): """The shared CuteDSL MoE per-forward token bound. Fields are set directly to exercise the math independently of __post_init__ resolution.