set default value as false

Author: gongshaotian

fastdeploy/config.py

@@ -592,7 +592,7 @@ def __init__(
         self.use_unique_memory_pool: bool = False
 
         """ Whether to use cudagraph for draft model."""
-        self.draft_model_use_cudagraph: bool = True
+        self.draft_model_use_cudagraph: bool = False
 
         self.max_capture_size: int = None
         self.real_shape_to_captured_size: dict[int, int] = None

fastdeploy/worker/gpu_model_runner.py

@@ -1371,31 +1371,35 @@ def capture_model(self) -> None:
                         expected_decode_len=1,
                     )
                     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.parallel_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.parallel_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,
                     )
                     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.parallel_config.max_num_batched_tokens,
-                    batch_size=int(1),
-                    in_capturing=True,
-                    expected_decode_len=3,
-                    accept_all_drafts=False,
-                )
-                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):
                 self._dummy_run(