[Graph Optimization] Add the CUDAGraph usage switch for Draft Model

fastdeploy/config.py

@@ -823,6 +823,8 @@ def __init__(
         self.real_shape_to_captured_size: dict[int, int] = None
         """ Whether to use shared memory pool for multi capture_size """
         self.use_unique_memory_pool: bool = True
+        """ Whether to use cudagraph for draft model."""
+        self.draft_model_use_cudagraph: bool = False
 
         # CINN Config ...
         if args is not None:

fastdeploy/spec_decode/mtp.py

@@ -83,6 +83,7 @@ def __init__(
         self._init_model_inputs()
 
         # CUDA Graph
+        self.draft_model_use_cudagraph = self.graph_opt_config.draft_model_use_cudagraph
         self.cudagraph_capture_sizes = list(reversed(self.graph_opt_config.cudagraph_capture_sizes))
         self.sot_warmup_sizes = self.graph_opt_config.sot_warmup_sizes
 
@@ -624,7 +625,7 @@ def _initialize_forward_meta(self, step_use_cudagraph: bool = False):
         for attn_backend in self.attn_backends:
             attn_backend.init_attention_metadata(self.forward_meta)
 
-        self.forward_meta.step_use_cudagraph = step_use_cudagraph
+        self.forward_meta.step_use_cudagraph = step_use_cudagraph and self.draft_model_use_cudagraph
 
     def exist_prefill(self):
         """

fastdeploy/worker/gpu_model_runner.py

@@ -1884,34 +1884,37 @@ def capture_model(self) -> None:
                         logger.info(
                             f"Warm up the Target model with the num_tokens:{batch_size}, expected_decode_len:{1}"
                         )
-                # Capture Draft Model without bsz 1
-                # NOTE(liujundong): expected_decode_len = 1, will affect mtp capture in cudagraph
-                for batch_size in sorted(capture_sizes, reverse=True):
-                    if batch_size == 1:
-                        logger.info("Skip token_num = 1, when capture Draft model for mtp")
-                    else:
-                        assert batch_size % 2 == 0
+                if self.graph_opt_config.draft_model_use_cudagraph:
+                    # Capture Draft Model without bsz 1
+                    # NOTE(liujundong): expected_decode_len = 1, will affect mtp capture in cudagraph
+                    for batch_size in sorted(capture_sizes, reverse=True):
+                        if batch_size == 1:
+                            logger.info("Skip token_num = 1, when capture Draft model for mtp")
+                        else:
+                            assert batch_size % 2 == 0
+                            self._dummy_run(
+                                num_tokens=self.scheduler_config.max_num_batched_tokens,
+                                batch_size=int(batch_size / 2),
+                                in_capturing=True,
+                                expected_decode_len=3,
+                                accept_all_drafts=True,
+                            )
+                            logger.info(
+                                f"Warm up the Draft model with the num_tokens:{batch_size}, expected_decode_len:{3}"
+                            )
+                    # Capture Draft Model with bsz 1
+                    if 1 in capture_sizes:
                         self._dummy_run(
                             num_tokens=self.scheduler_config.max_num_batched_tokens,
-                            batch_size=int(batch_size / 2),
+                            batch_size=int(1),
                             in_capturing=True,
                             expected_decode_len=3,
-                            accept_all_drafts=True,
+                            accept_all_drafts=False,
+                            reject_all_drafts=True,
                         )
                         logger.info(
                             f"Warm up the Draft model with the num_tokens:{batch_size}, expected_decode_len:{3}"
                         )
-                # Capture Draft Model with bsz 1
-                if 1 in capture_sizes:
-                    self._dummy_run(
-                        num_tokens=self.scheduler_config.max_num_batched_tokens,
-                        batch_size=int(1),
-                        in_capturing=True,
-                        expected_decode_len=3,
-                        accept_all_drafts=False,
-                        reject_all_drafts=True,
-                    )
-                    logger.info(f"Warm up the Draft model with the num_tokens:{batch_size}, expected_decode_len:{3}")
 
             else:
                 for batch_size in sorted(capture_sizes, reverse=True):

tests/e2e/test_ernie_21b_mtp.py

@@ -129,6 +129,8 @@ def setup_and_run_server():
         "wint4",
         "--speculative-config",
         json.dumps(speculative_config),
+        "--graph-optimization-config",
+        '{"use_cudagraph":true,  "use_unique_memory_pool":true, "draft_model_use_cudagraph":true}',
     ]
 
     # Start subprocess in new process group