From 902a33047496ca7d2e7af3261f6fc082504b471a Mon Sep 17 00:00:00 2001 From: gongshaotian Date: Mon, 27 Oct 2025 14:56:34 +0800 Subject: [PATCH 1/3] add draft model using cudagraph switch --- fastdeploy/config.py | 2 ++ fastdeploy/spec_decode/mtp.py | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/fastdeploy/config.py b/fastdeploy/config.py index 8cd82fbc655..46199507e96 100644 --- a/fastdeploy/config.py +++ b/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 = True # CINN Config ... if args is not None: diff --git a/fastdeploy/spec_decode/mtp.py b/fastdeploy/spec_decode/mtp.py index 70127c8eb29..fedbd52a883 100644 --- a/fastdeploy/spec_decode/mtp.py +++ b/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): """ From e5adf59e992c689137215aa82a4da3843c3052d4 Mon Sep 17 00:00:00 2001 From: gongshaotian Date: Tue, 28 Oct 2025 18:18:05 +0800 Subject: [PATCH 2/3] set default as false --- fastdeploy/config.py | 2 +- fastdeploy/worker/gpu_model_runner.py | 43 ++++++++++++++------------- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/fastdeploy/config.py b/fastdeploy/config.py index 46199507e96..21af68219f9 100644 --- a/fastdeploy/config.py +++ b/fastdeploy/config.py @@ -824,7 +824,7 @@ def __init__( """ 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 = True + self.draft_model_use_cudagraph: bool = False # CINN Config ... if args is not None: diff --git a/fastdeploy/worker/gpu_model_runner.py b/fastdeploy/worker/gpu_model_runner.py index 598388bf9dc..12b35ec71b1 100644 --- a/fastdeploy/worker/gpu_model_runner.py +++ b/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): From cb364443ef7fdde989d971013827bf7fc5c20d91 Mon Sep 17 00:00:00 2001 From: gongshaotian Date: Tue, 28 Oct 2025 21:45:06 +0800 Subject: [PATCH 3/3] capture draft model in ci --- tests/e2e/test_ernie_21b_mtp.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/test_ernie_21b_mtp.py b/tests/e2e/test_ernie_21b_mtp.py index bff9d7ebcb6..d9bc778ede0 100644 --- a/tests/e2e/test_ernie_21b_mtp.py +++ b/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