support spec with tp and cudagraph

Author: RunningLeon

benchmark/benchmark_serving.py

@@ -13,14 +13,15 @@ def get_launching_server_cmd(model_path, backend, server_config):
     elif backend == 'sglang':
         cmd = ['python3', '-m', 'sglang.launch_server', '--model-path', model_path]
     elif backend == 'vllm':
-        cmd = ['vllm', 'serve', '--model', model_path]
+        cmd = ['vllm', 'serve', model_path]
     else:
         raise ValueError(f'unknown backend: {backend}')
     for key, value in server_config.items():
         # Convert snake_case to kebab-case for command line args
         key = key.replace('_', '-')
         cmd.append(f'--{key}')
-        cmd.append(str(value))
+        if str(value):
+            cmd.append(str(value))
     # Special handling for proxy server case
     if server_config.get('proxy_url') and server_config.get('dp'):
         cmd.append('--allow-terminate-by-client')
@@ -66,9 +67,9 @@ def get_server_ip_port(backend: str, server_config: Dict) -> Tuple[str, int]:
             server_ip = server_config.get('server_ip', '0.0.0.0')
             server_port = server_config.get('server_port', 23333)
     elif backend == 'sglang':
-        return (server_config.get('server_ip', '0.0.0.0'), server_config.get('server_port', 30000))
+        return (server_config.get('server_ip', '0.0.0.0'), server_config.get('port', 30000))
     elif backend == 'vllm':
-        return (server_config.get('server_ip', '0.0.0.0'), server_config.get('server_port', 8000))
+        return (server_config.get('server_ip', '0.0.0.0'), server_config.get('port', 8000))
     else:
         raise ValueError(f'unknown backend: {backend}')
     return server_ip, server_port
@@ -131,7 +132,7 @@ def benchmark(model_path: str, backend: str, server_config: Dict, data_config: D
 
     try:
 
-        print(f"Starting api_server: {' '.join(server_cmd)}")
+        print(f"Starting api_server: {' '.join(server_cmd)}", flush=True)
         proc = subprocess.Popen(server_cmd)
         # Wait for the server to be ready
         wait_server_ready(server_ip, server_port)

lmdeploy/api.py

@@ -3,7 +3,7 @@
 from typing import List, Literal, Optional, Union
 
 from .archs import autoget_backend_config, get_task
-from .messages import PytorchEngineConfig, TurbomindEngineConfig
+from .messages import PytorchEngineConfig, SpeculativeConfig, TurbomindEngineConfig
 from .model import ChatTemplateConfig
 
 
@@ -12,6 +12,7 @@ def pipeline(model_path: str,
              chat_template_config: Optional[ChatTemplateConfig] = None,
              log_level: str = 'WARNING',
              max_log_len: int = None,
+             speculative_config: SpeculativeConfig = None,
              **kwargs):
     """
     Args:
@@ -68,6 +69,12 @@ def pipeline(model_path: str,
             if backend_config is not None else None
         model_path = get_model(model_path, download_dir, revision)
 
+    # spec model
+    if speculative_config is not None and speculative_config.model and not os.path.exists(speculative_config.model):
+        download_dir = backend_config.download_dir \
+            if backend_config is not None else None
+        speculative_config.model = get_model(speculative_config.model, download_dir)
+
     _, pipeline_class = get_task(model_path)
     if not isinstance(backend_config, PytorchEngineConfig):
         # set auto backend mode
@@ -80,6 +87,7 @@ def pipeline(model_path: str,
                           backend_config=backend_config,
                           chat_template_config=chat_template_config,
                           max_log_len=max_log_len,
+                          speculative_config=speculative_config,
                           **kwargs)
 
 

lmdeploy/cli/cli.py

@@ -3,7 +3,8 @@
 import os
 
 from ..version import __version__
-from .utils import ArgumentHelper, DefaultsAndTypesHelpFormatter, FlexibleArgumentParser, convert_args
+from .utils import (ArgumentHelper, DefaultsAndTypesHelpFormatter, FlexibleArgumentParser, convert_args,
+                    get_speculative_config)
 
 
 class CLI(object):
@@ -44,12 +45,13 @@ def add_parser_chat():
                             ', "baichuan-inc/baichuan2-7b-chat" and so on')
         # common args
         ArgumentHelper.backend(parser)
-        # # chat template args
+        ArgumentHelper.log_level(parser)
+        # chat template args
         ArgumentHelper.chat_template(parser)
         # model args
         ArgumentHelper.revision(parser)
         ArgumentHelper.download_dir(parser)
-        #
+
         # pytorch engine args
         pt_group = parser.add_argument_group('PyTorch engine arguments')
         ArgumentHelper.adapters(pt_group)
@@ -77,6 +79,9 @@ def add_parser_chat():
         ArgumentHelper.rope_scaling_factor(tb_group)
         ArgumentHelper.communicator(tb_group)
 
+        # speculative decoding
+        ArgumentHelper.add_spec_group(parser)
+
     @staticmethod
     def add_parser_checkenv():
         """Add parser for check_env command."""
@@ -168,7 +173,13 @@ def get_gpu_topo():
     @staticmethod
     def chat(args):
         from .chat import main
+
         kwargs = convert_args(args)
+        speculative_config = get_speculative_config(args)
+        to_remove = ['speculative_algorithm', 'speculative_draft_model', 'speculative_num_draft_tokens']
+        for key in to_remove:
+            kwargs.pop(key)
+        kwargs['speculative_config'] = speculative_config
         main(**kwargs)
 
     @staticmethod

lmdeploy/cli/serve.py

@@ -3,7 +3,8 @@
 from lmdeploy.utils import get_max_batch_size
 
 from .cli import CLI
-from .utils import ArgumentHelper, DefaultsAndTypesHelpFormatter, convert_args, get_chat_template, get_lora_adapters
+from .utils import (ArgumentHelper, DefaultsAndTypesHelpFormatter, convert_args, get_chat_template, get_lora_adapters,
+                    get_speculative_config)
 
 
 class SubCliServe:
@@ -144,6 +145,9 @@ def add_parser_api_server():
         vision_group = parser.add_argument_group('Vision model arguments')
         ArgumentHelper.vision_max_batch_size(vision_group)
 
+        # spec decode
+        ArgumentHelper.add_spec_group(parser)
+
     @staticmethod
     def add_parser_proxy():
         """Add parser for proxy server command."""
@@ -247,61 +251,68 @@ def api_server(args):
                                                    enable_metrics=args.enable_metrics,
                                                    hf_overrides=args.hf_overrides)
         chat_template_config = get_chat_template(args.chat_template, args.model_path)
+        speculative_config = get_speculative_config(args)
 
         from lmdeploy.messages import VisionConfig
         vision_config = VisionConfig(args.vision_max_batch_size)
         if args.dp == 1:
             from lmdeploy.serve.openai.api_server import serve as run_api_server
 
-            run_api_server(args.model_path,
-                           model_name=args.model_name,
-                           backend=backend,
-                           backend_config=backend_config,
-                           chat_template_config=chat_template_config,
-                           vision_config=vision_config,
-                           server_name=args.server_name,
-                           server_port=args.server_port,
-                           allow_origins=args.allow_origins,
-                           allow_credentials=args.allow_credentials,
-                           allow_methods=args.allow_methods,
-                           allow_headers=args.allow_headers,
-                           allow_terminate_by_client=args.allow_terminate_by_client,
-                           log_level=args.log_level.upper(),
-                           api_keys=args.api_keys,
-                           ssl=args.ssl,
-                           proxy_url=args.proxy_url,
-                           max_log_len=args.max_log_len,
-                           disable_fastapi_docs=args.disable_fastapi_docs,
-                           max_concurrent_requests=args.max_concurrent_requests,
-                           reasoning_parser=args.reasoning_parser,
-                           tool_call_parser=args.tool_call_parser)
+            run_api_server(
+                args.model_path,
+                model_name=args.model_name,
+                backend=backend,
+                backend_config=backend_config,
+                chat_template_config=chat_template_config,
+                vision_config=vision_config,
+                server_name=args.server_name,
+                server_port=args.server_port,
+                allow_origins=args.allow_origins,
+                allow_credentials=args.allow_credentials,
+                allow_methods=args.allow_methods,
+                allow_headers=args.allow_headers,
+                allow_terminate_by_client=args.allow_terminate_by_client,
+                log_level=args.log_level.upper(),
+                api_keys=args.api_keys,
+                ssl=args.ssl,
+                proxy_url=args.proxy_url,
+                max_log_len=args.max_log_len,
+                disable_fastapi_docs=args.disable_fastapi_docs,
+                max_concurrent_requests=args.max_concurrent_requests,
+                reasoning_parser=args.reasoning_parser,
+                tool_call_parser=args.tool_call_parser,
+                speculative_config=speculative_config,
+            )
         else:
             from lmdeploy.serve.openai.launch_server import launch_server
 
-            launch_server(args.nnodes,
-                          args.node_rank,
-                          args.model_path,
-                          model_name=args.model_name,
-                          backend=backend,
-                          backend_config=backend_config,
-                          chat_template_config=chat_template_config,
-                          vision_config=vision_config,
-                          server_name=args.server_name,
-                          server_port=args.server_port,
-                          allow_origins=args.allow_origins,
-                          allow_credentials=args.allow_credentials,
-                          allow_methods=args.allow_methods,
-                          allow_headers=args.allow_headers,
-                          allow_terminate_by_client=args.allow_terminate_by_client,
-                          log_level=args.log_level.upper(),
-                          api_keys=args.api_keys,
-                          ssl=args.ssl,
-                          proxy_url=args.proxy_url,
-                          max_log_len=args.max_log_len,
-                          disable_fastapi_docs=args.disable_fastapi_docs,
-                          max_concurrent_requests=args.max_concurrent_requests,
-                          reasoning_parser=args.reasoning_parser,
-                          tool_call_parser=args.tool_call_parser)
+            launch_server(
+                args.nnodes,
+                args.node_rank,
+                args.model_path,
+                model_name=args.model_name,
+                backend=backend,
+                backend_config=backend_config,
+                chat_template_config=chat_template_config,
+                vision_config=vision_config,
+                server_name=args.server_name,
+                server_port=args.server_port,
+                allow_origins=args.allow_origins,
+                allow_credentials=args.allow_credentials,
+                allow_methods=args.allow_methods,
+                allow_headers=args.allow_headers,
+                allow_terminate_by_client=args.allow_terminate_by_client,
+                log_level=args.log_level.upper(),
+                api_keys=args.api_keys,
+                ssl=args.ssl,
+                proxy_url=args.proxy_url,
+                max_log_len=args.max_log_len,
+                disable_fastapi_docs=args.disable_fastapi_docs,
+                max_concurrent_requests=args.max_concurrent_requests,
+                reasoning_parser=args.reasoning_parser,
+                tool_call_parser=args.tool_call_parser,
+                speculative_config=speculative_config,
+            )
 
     @staticmethod
     def proxy(args):

lmdeploy/cli/utils.py

@@ -99,6 +99,19 @@ def get_chat_template(chat_template: str, model_path: str = None):
         return None
 
 
+def get_speculative_config(args):
+    """Get speculative config from args."""
+    from lmdeploy.messages import SpeculativeConfig
+    speculative_config = None
+    if args.speculative_algorithm is not None:
+        speculative_config = SpeculativeConfig(
+            method=args.speculative_algorithm,
+            model=args.speculative_draft_model,
+            num_speculative_tokens=args.speculative_num_draft_tokens,
+        )
+    return speculative_config
+
+
 class ArgumentHelper:
     """Helper class to add unified argument."""
 
@@ -654,6 +667,26 @@ def dllm_confidence_threshold(parser):
                                    default=0.85,
                                    help='The confidence threshold for dllm.')
 
+    def add_spec_group(parser):
+        spec_group = parser.add_argument_group('Speculative decoding arguments')
+        spec_group.add_argument('--speculative-algorithm',
+                                type=str,
+                                default=None,
+                                choices=['eagle', 'eagle3', 'deepseek_mtp'],
+                                help='The speculative algorithm to use. `None` means speculative decoding is disabled')
+
+        spec_group.add_argument('--speculative-draft-model',
+                                type=str,
+                                default=None,
+                                help='The path to speculative draft model')
+
+        spec_group.add_argument('--speculative-num-draft-tokens',
+                                type=int,
+                                default=1,
+                                help='The number of speculative tokens to generate per step')
+
+        return spec_group
+
 
 # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/utils/__init__.py
 class FlexibleArgumentParser(argparse.ArgumentParser):

lmdeploy/messages.py

@@ -522,6 +522,7 @@ class RequestMetrics:
     """
     token_timestamp: float = 0.0
     engine_events: List[EngineEvent] = field(default_factory=list)
+    spec_info: Optional[Dict[str, Any]] = None
 
 
 @dataclass
@@ -562,3 +563,17 @@ class VisionConfig:
     """
     max_batch_size: int = 1
     thread_safe: bool = False
+
+
+@dataclass
+class SpeculativeConfig:
+    """Speculative decoding config.
+
+    Args:
+        method (str): the speculative decoding method.
+        model (str): the path of speculative model.
+        num_speculative_tokens (int): number of generated token of draft model per step
+    """
+    method: str
+    model: str = ''
+    num_speculative_tokens: int = 1