[Refactor] Cuda Graph Runner/Backend Refactor (#23906)

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

test/registered/ascend/basic_function/optimization_debug/test_npu_piecewise_graph_prefill.py

@@ -29,8 +29,8 @@ class TestPiecewiseGraphPrefillCorrectness(GSM8KAscendMixin, CustomTestCase):
         "ascend",
         "--cuda-graph-bs",
         128,
-        "--enforce-piecewise-cuda-graph",
-        "--piecewise-cuda-graph-tokens",
+        "--cuda-graph-backend-prefill=tc_piecewise",
+        "--cuda-graph-bs-prefill",
         *TOKENS_TO_CAPTURE,
     ]
     accuracy = 0.84
@@ -45,8 +45,8 @@ class TestPiecewiseGraphPrefillBenchmark(CustomTestCase):
         0.8,
         "--attention-backend",
         "ascend",
-        "--enforce-piecewise-cuda-graph",
-        "--piecewise-cuda-graph-tokens",
+        "--cuda-graph-backend-prefill=tc_piecewise",
+        "--cuda-graph-bs-prefill",
     ] + TOKENS_TO_CAPTURE
 
     latency = 0.045

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -394,7 +394,15 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
         )
         from sglang.srt.server_args import ServerArgs
 
-        self.assertFalse(ServerArgs(model_path="dummy").enable_breakable_cuda_graph)
+        # cg-refactor folded the legacy enable_breakable_cuda_graph flag
+        # into cuda_graph_config. Verify the per-phase backend selectors
+        # default to None (i.e. nothing opted into BREAKABLE without an
+        # explicit CLI flag).
+        sa = ServerArgs(model_path="dummy")
+        self.assertNotEqual(sa.cuda_graph_backend_decode, "breakable")
+        self.assertNotEqual(sa.cuda_graph_backend_prefill, "breakable")
+        self.assertNotEqual(sa.decode_cuda_graph_backend, "breakable")
+        self.assertNotEqual(sa.prefill_cuda_graph_backend, "breakable")
         self.assertFalse(
             AttentionBackend.use_captured_forward_metadata_for_breakable_cuda_graph
         )

test/registered/core/test_no_extra_forked_cuda_context.py

@@ -34,8 +34,10 @@ class TestTPServerGPUProcesses(CustomTestCase):
                 str(cls.tp_size),
                 "--mem-fraction-static",
                 "0.70",
-                "--disable-cuda-graph",
-                "--disable-piecewise-cuda-graph",
+                "--cuda-graph-backend-decode",
+                "disabled",
+                "--cuda-graph-backend-prefill",
+                "disabled",
             ],
         )
 

test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py

@@ -1,10 +1,10 @@
 """Tests for the breakable CUDA graph (BCG) runner.
 
 Two test classes:
-- ``TestBreakableCUDAGraphBasic`` / ``TestCopyOutput`` / ``TestBreakGraphHelper``:
+- TestBreakableCUDAGraphBasic / TestCopyOutput / TestBreakGraphHelper:
   unit tests for the core capture / replay mechanism (simple tensor ops).
-- ``TestBreakableCudaGraph``: integration test — spin up Qwen3-8B with
-  ``--enable-breakable-cuda-graph`` and check mgsm_en accuracy.
+- TestBreakableCudaGraph: integration test — spin up Qwen3-8B with
+  --enable-breakable-cuda-graph and check mgsm_en accuracy.
 """
 
 import unittest
@@ -53,7 +53,7 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         except ImportError:
             raise unittest.SkipTest("cuda-python not installed")
 
-        from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
             BreakableCUDAGraph,
             BreakableCUDAGraphCapture,
             eager_on_graph,
@@ -199,7 +199,7 @@ class TestCopyOutput(CustomTestCase):
         except ImportError:
             raise unittest.SkipTest("cuda-python not installed")
 
-        from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
             _copy_output,
         )
 
@@ -261,7 +261,7 @@ class TestBreakGraphHelper(CustomTestCase):
         except ImportError:
             raise unittest.SkipTest("cuda-python not installed")
 
-        from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
             BreakableCUDAGraph,
             BreakableCUDAGraphCapture,
             break_graph,
@@ -303,7 +303,7 @@ class TestBreakableCudaGraph(CustomTestCase):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
-                "--enable-breakable-cuda-graph",
+                "--cuda-graph-backend-prefill=breakable",
             ],
         )
 

test/registered/cuda_graph/piecewise/test_pcg_glm5_fp4.py

@@ -43,7 +43,7 @@ class TestPCGGlm5Fp4(CustomTestCase):
                 "--quantization",
                 "modelopt_fp4",
                 "--disable-flashinfer-autotune",
-                "--enforce-piecewise-cuda-graph",
+                "--cuda-graph-backend-prefill=tc_piecewise",
                 "--model-loader-extra-config",
                 '{"enable_multithread_load": true, "num_threads": 64}',
             ],

test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding.py

@@ -22,7 +22,7 @@ class TestPCGWithEAGLE3(PCGSpecBase, unittest.TestCase):
         "--tp",
         "2",
         "--trust-remote-code",
-        "--enforce-piecewise-cuda-graph",
+        "--cuda-graph-backend-prefill=tc_piecewise",
         "--mem-fraction-static",
         "0.6",
         "--speculative-algorithm",

test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py

@@ -26,7 +26,8 @@ class TestPCGWithDFlash(PCGSpecBase, CustomTestCase):
         "--trust-remote-code",
         "--attention-backend",
         "flashinfer",
-        "--enforce-piecewise-cuda-graph",
+        "--cuda-graph-backend-prefill",
+        "tc_piecewise",
         "--speculative-algorithm",
         "DFLASH",
         "--speculative-draft-model-path",
@@ -35,7 +36,7 @@ class TestPCGWithDFlash(PCGSpecBase, CustomTestCase):
         "1",
         "--max-running-requests",
         "64",
-        "--cuda-graph-bs",
+        "--cuda-graph-bs-decode",
         *[str(i) for i in range(1, 65)],
     ]
     server_env = {"SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1"}

test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_extra.py

@@ -24,7 +24,6 @@ class TestPCGWithMTP(PCGSpecBase, unittest.TestCase):
         "fp8",
         "--mamba-scheduler-strategy",
         "extra_buffer",
-        "--enable-piecewise-cuda-graph",
         "--speculative-algorithm",
         "NEXTN",
         "--reasoning-parser",
@@ -42,7 +41,7 @@ class TestPCGWithSTANDALONE(PCGSpecBase, unittest.TestCase):
     model = "meta-llama/Llama-3.1-8B-Instruct"
     server_args = [
         "--trust-remote-code",
-        "--enforce-piecewise-cuda-graph",
+        "--cuda-graph-backend-prefill=tc_piecewise",
         "--mem-fraction-static",
         "0.5",
         "--speculative-algorithm",
@@ -65,7 +64,7 @@ class TestPCGWithNGRAM(PCGSpecBase, unittest.TestCase):
     model = "Qwen/Qwen2.5-Coder-7B-Instruct"
     server_args = [
         "--trust-remote-code",
-        "--enforce-piecewise-cuda-graph",
+        "--cuda-graph-backend-prefill=tc_piecewise",
         "--speculative-algorithm",
         "NGRAM",
         "--speculative-num-draft-tokens",

test/registered/cuda_graph/piecewise/test_piecewise_cuda_graph_support_1_gpu.py

@@ -33,7 +33,7 @@ class TestPiecewiseCudaGraphQwen25VL(CustomTestCase):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
-                "--enforce-piecewise-cuda-graph",
+                "--cuda-graph-backend-prefill=tc_piecewise",
                 "--disable-radix-cache",
             ],
         )
@@ -69,7 +69,7 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
             model_path=model_path,
             enable_multimodal=True,
             is_embedding=True,
-            enforce_piecewise_cuda_graph=True,
+            cuda_graph_backend_prefill="tc_piecewise",
         )
         out = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0]["embedding"]
         engine.shutdown()
@@ -79,7 +79,7 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase):
             model_path=model_path,
             enable_multimodal=True,
             is_embedding=True,
-            disable_piecewise_cuda_graph=True,
+            cuda_graph_backend_prefill="disabled",
         )
         out_without_pcg = engine.encode([text], image_data=[DEFAULT_IMAGE_URL])[0][
             "embedding"

test/registered/debug_utils/test_engine_dumper_comparator_e2e.py

@@ -325,7 +325,7 @@ def _run_server_and_generate(
         "--mem-fraction-static",
         "0.5",
         "--disable-cuda-graph",
-        "--disable-piecewise-cuda-graph",
+        "--cuda-graph-backend-prefill=disabled",
         "--disable-radix-cache",
     ]
     if extra_server_args:

test/registered/kv_canary/test_self_unit_capacities.py

@@ -4,17 +4,24 @@ import unittest
 from types import SimpleNamespace
 
 from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    CudaGraphConfig,
+    PhaseConfig,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small")
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
 class TestComputeLaunchCapacities(CustomTestCase):
     @staticmethod
     def _make_server_args(*, max_bs: int) -> SimpleNamespace:
         return SimpleNamespace(
-            cuda_graph_max_bs=max_bs,
+            cuda_graph_config=CudaGraphConfig(
+                decode=PhaseConfig(backend=Backend.FULL, max_bs=max_bs)
+            ),
             speculative_num_draft_tokens=0,
             chunked_prefill_size=None,
             max_prefill_tokens=128,

test/registered/lora/test_lora_update.py

@@ -879,7 +879,7 @@ class LoRAUpdateTestSessionBase:
         lora_target_modules: Optional[List[str]] = None,
         lora_backend: str = "csgmv",
         disable_cuda_graph: bool = False,
-        cuda_graph_max_bs: int = 4,
+        cuda_graph_max_bs_decode: int = 4,
     ):
         self.testcase = testcase
         self.model_path = model_path
@@ -890,7 +890,7 @@ class LoRAUpdateTestSessionBase:
         self.max_loaded_loras = max_loaded_loras
         self.lora_backend = lora_backend
         self.disable_cuda_graph = disable_cuda_graph
-        self.cuda_graph_max_bs = cuda_graph_max_bs
+        self.cuda_graph_max_bs_decode = cuda_graph_max_bs_decode
         self.enable_lora = enable_lora
 
         self.expected_adapters = set()
@@ -964,7 +964,7 @@ class LoRAUpdateEngineTestSession(LoRAUpdateTestSessionBase):
             max_loras_per_batch=self.max_loras_per_batch,
             max_loaded_loras=self.max_loaded_loras,
             disable_cuda_graph=self.disable_cuda_graph,
-            cuda_graph_max_bs=self.cuda_graph_max_bs,
+            cuda_graph_max_bs_decode=self.cuda_graph_max_bs_decode,
             enable_lora=self.enable_lora,
             disable_radix_cache=True,
         )
@@ -1101,8 +1101,8 @@ class LoRAUpdateServerTestSession(LoRAUpdateTestSessionBase):
 
     def __enter__(self):
         other_args = [
-            "--cuda-graph-max-bs",
-            str(self.cuda_graph_max_bs),
+            "--cuda-graph-max-bs-decode",
+            str(self.cuda_graph_max_bs_decode),
             "--max-loras-per-batch",
             str(self.max_loras_per_batch),
             "--lora-backend",

test/registered/model_loading/test_external_models.py

@@ -23,7 +23,7 @@ class TestExternalModels(CustomTestCase):
 
         engine = sgl.Engine(
             model_path=model_path,
-            cuda_graph_max_bs=1,
+            cuda_graph_max_bs_decode=1,
             max_total_tokens=64,
             enable_multimodal=True,
         )

test/registered/model_loading/test_load_weights_from_remote_instance.py

@@ -188,7 +188,7 @@ def init_process_dst(
             model_path=model_name,
             base_gpu_id=base_gpu_id,
             tp_size=tp_size,
-            cuda_graph_max_bs=2,
+            cuda_graph_max_bs_decode=2,
             tokenizer_path=model_name,
             remote_instance_weight_loader_seed_instance_ip=seed_instance_ip,
             remote_instance_weight_loader_seed_instance_service_port=seed_instance_service_port,

test/registered/model_loading/test_load_weights_from_remote_instance_npu.py

@@ -197,7 +197,7 @@ def init_process_dst(
             model_path=model_name,
             base_gpu_id=base_gpu_id,
             tp_size=tp_size,
-            cuda_graph_max_bs=2,
+            cuda_graph_max_bs_decode=2,
             tokenizer_path=model_name,
             remote_instance_weight_loader_seed_instance_ip=seed_instance_ip,
             remote_instance_weight_loader_seed_instance_service_port=seed_instance_service_port,

test/registered/model_loading/test_prefetch_checkpoints_multi_gpu.py

@@ -26,7 +26,7 @@ class TestPrefetchCheckpointsMultiGPU(CustomTestCase):
             enable_dp_attention=True,
             disable_radix_cache=True,
             weight_loader_prefetch_checkpoints=True,
-            cuda_graph_max_bs=1,
+            cuda_graph_max_bs_decode=1,
             max_total_tokens=256,
         )
 

test/registered/model_loading/test_runai_model_loader.py

@@ -28,7 +28,7 @@ class TestRunaiModelLoader(CustomTestCase):
             cls.engine = sgl.Engine(
                 model_path=TEST_GCS_MODEL,
                 load_format="runai_streamer",
-                cuda_graph_max_bs=1,
+                cuda_graph_max_bs_decode=1,
                 max_total_tokens=64,
             )
 

test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

@@ -179,8 +179,8 @@ class TestDSV4FlashFP4BreakableCudaGraphB200(
                 "4",
                 "--enable-dp-attention",
                 "--enable-mixed-chunk",
-                "--enable-breakable-cuda-graph",
-                "--enforce-piecewise-cuda-graph",
+                "--cuda-graph-backend-prefill",
+                "breakable",
                 "--moe-a2a-backend",
                 "deepep",
                 "--deepep-config",

test/registered/quant/test_fp8_gemm_sm120.py

@@ -33,7 +33,7 @@ class FP8GemmSM120Base:
             "--trust-remote-code",
             "--fp8-gemm-backend",
             cls.backend,
-            "--disable-piecewise-cuda-graph",
+            "--cuda-graph-backend-prefill=disabled",
         ]
         if cls.quantization:
             other_args += ["--quantization", cls.quantization]

test/registered/quant/test_gguf.py

@@ -19,7 +19,9 @@ class TestGGUF(CustomTestCase):
             filename="qwen2-1_5b-instruct-q4_k_m.gguf",
         )
 
-        engine = sgl.Engine(model_path=model_path, random_seed=42, cuda_graph_max_bs=2)
+        engine = sgl.Engine(
+            model_path=model_path, random_seed=42, cuda_graph_max_bs_decode=2
+        )
         outputs = engine.generate(prompt, sampling_params)["text"]
         engine.shutdown()
 

test/registered/quant/test_nvfp4_gemm_sm120.py

@@ -32,7 +32,7 @@ class FP4GemmSM120Base:
             "modelopt_fp4",
             "--fp4-gemm-backend",
             cls.backend,
-            "--disable-piecewise-cuda-graph",
+            "--cuda-graph-backend-prefill=disabled",
         ]
         cls.process = popen_launch_server(
             cls.model,

test/registered/radix_cache/test_swa_radix_cache_kl.py

@@ -18,7 +18,7 @@ class TestSWARadixCacheKL(KLDivergenceMixin, DefaultServerBase):
         "1",
         "--mem-fraction-static",
         "0.70",
-        "--disable-piecewise-cuda-graph",
+        "--cuda-graph-backend-prefill=disabled",
     ]
 
 

test/registered/rl/test_update_weights_from_distributed.py

@@ -315,7 +315,7 @@ def init_process_sgl(
             model_path=model_name,
             base_gpu_id=base_gpu_id,
             tp_size=tp_size,
-            cuda_graph_max_bs=2,
+            cuda_graph_max_bs_decode=2,
         )
     else:
         if rank == 1:

test/registered/sessions/test_session_control.py

@@ -45,7 +45,7 @@ class TestSessionControl(CustomTestCase):
                 "--attention-backend",
                 "triton",
                 "--disable-cuda-graph",
-                "--disable-piecewise-cuda-graph",
+                "--cuda-graph-backend-prefill=disabled",
             ],
         )
 

test/registered/sessions/test_session_latency.py

@@ -271,7 +271,7 @@ class TestSessionLatency(CustomTestCase):
                 "--enable-streaming-session",
                 "--mem-fraction-static",
                 "0.70",
-                "--disable-piecewise-cuda-graph",
+                "--cuda-graph-backend-prefill=disabled",
                 "--page-size",
                 "4",
             ],

test/registered/sessions/test_streaming_session_swa.py

@@ -21,7 +21,7 @@ SWA_MODEL = "openai/gpt-oss-20b"
 SWA_COMMON_ARGS = [
     "--mem-fraction-static",
     "0.70",
-    "--disable-piecewise-cuda-graph",
+    "--cuda-graph-backend-prefill=disabled",
 ]
 
 

test/registered/unit/hardware_backend/mlx/test_attention_patching.py

@@ -396,7 +396,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
         )
         scheduler.server_args = SimpleNamespace(
             enable_two_batch_overlap=False,
-            disable_piecewise_cuda_graph=True,
+            cuda_graph_config=None,
         )
         scheduler.spec_algorithm = SpeculativeAlgorithm.NONE
         scheduler.req_to_token_pool = ReqToTokenPool(

test/registered/unit/managers/test_customized_info_streaming.py

@@ -93,7 +93,6 @@ class TestCustomizedInfoStreaming(CustomTestCase):
             incremental_streaming_output=True,
             skip_tokenizer_init=True,
             disable_cuda_graph=True,
-            disable_piecewise_cuda_graph=True,
             disable_radix_cache=True,
             random_seed=0,
             log_level="error",

test/registered/unit/server_args/test_server_args.py

@@ -3,10 +3,16 @@ import json
 import os
 import tempfile
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import sglang.srt.server_args as server_args_module
 from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    CudaGraphConfig,
+    PhaseConfig,
+)
 from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import (
@@ -621,9 +627,61 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
             ServerArgs(**self._base_kwargs(kv_cache_dtype="fp4_e2m1"))
 
 
-class TestCutedslMoeMaxNumTokens(unittest.TestCase):
+class TestCudaGraphConfigDataclassAccess(CustomTestCase):
+    def test_overlap_force_cpu_seq_lens_with_tc_piecewise_prefill(self):
+        from sglang.srt.managers.overlap_utils import decide_needs_cpu_seq_lens
+
+        server_args = SimpleNamespace(
+            enable_two_batch_overlap=False,
+            cuda_graph_config=CudaGraphConfig(
+                prefill=PhaseConfig(backend=Backend.TC_PIECEWISE)
+            ),
+        )
+        attn_backend = SimpleNamespace(needs_cpu_seq_lens=False)
+
+        self.assertTrue(decide_needs_cpu_seq_lens(server_args, [attn_backend]))
+
+    @patch(
+        "sglang.srt.model_executor.runner_backend."
+        "tc_piecewise_cuda_graph_backend.get_moe_a2a_backend"
+    )
+    def test_tc_piecewise_build_config_reads_phase_config_dataclass(
+        self, mock_get_moe_a2a_backend
+    ):
+        from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import (
+            TcPiecewiseCudaGraphBackend,
+        )
+
+        mock_backend = mock_get_moe_a2a_backend.return_value
+        mock_backend.is_deepep.return_value = False
+        mock_backend.is_mooncake.return_value = False
+        server_args = SimpleNamespace(
+            cuda_graph_config=CudaGraphConfig(
+                prefill=PhaseConfig(
+                    backend=Backend.TC_PIECEWISE,
+                    bs=[32, 64],
+                    tc_compiler="eager",
+                )
+            ),
+            enable_torch_compile_debug_mode=False,
+        )
+
+        config = TcPiecewiseCudaGraphBackend.build_compilation_config(server_args)
+
+        self.assertEqual(config.get_capture_sizes(), [32, 64])
+        self.assertEqual(config.compiler, "eager")
+
+
+class TestCutedslMoeMaxNumTokens(CustomTestCase):
     """The shared CuteDSL MoE per-forward token bound. Fields are set directly
-    to exercise the math independently of __post_init__ resolution."""
+    to exercise the math independently of __post_init__ resolution.
+
+    cg-refactor: the legacy disable_piecewise_cuda_graph /
+    piecewise_cuda_graph_max_tokens / cuda_graph_max_bs fields were
+    consolidated into cuda_graph_config; the helper accepts the legacy
+    kwarg names for test readability and translates them to the per-phase
+    dataclasses.
+    """
 
     def _args(self, **overrides):
         server_args = ServerArgs(model_path="dummy")
@@ -636,8 +694,21 @@ class TestCutedslMoeMaxNumTokens(unittest.TestCase):
             cuda_graph_max_bs=512,
         )
         fields.update(overrides)
+        disable_piecewise = fields.pop("disable_piecewise_cuda_graph")
+        piecewise_max = fields.pop("piecewise_cuda_graph_max_tokens")
+        cg_max_bs = fields.pop("cuda_graph_max_bs")
         for key, value in fields.items():
             setattr(server_args, key, value)
+        server_args.cuda_graph_config = CudaGraphConfig(
+            decode=PhaseConfig(backend=Backend.FULL, max_bs=cg_max_bs),
+            prefill=PhaseConfig(
+                backend=(
+                    Backend.DISABLED if disable_piecewise else Backend.TC_PIECEWISE
+                ),
+                max_bs=piecewise_max,
+                tc_compiler="eager",
+            ),
+        )
         return server_args
 
     def test_prefill_dominates_in_default_config(self):

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -50,7 +50,10 @@ def _make_worker(num_steps: int, num_draft_tokens: int):
     worker.device = DEVICE
     worker.speculative_num_steps = num_steps
     worker.speculative_num_draft_tokens = num_draft_tokens
-    worker.server_args = SimpleNamespace(cuda_graph_max_bs=8, max_running_requests=8)
+    worker.server_args = SimpleNamespace(
+        cuda_graph_config=SimpleNamespace(decode=SimpleNamespace(max_bs=8)),
+        max_running_requests=8,
+    )
     return worker