refactor(serve): simplify generation config merge and defaults

Drop PROTOCOL_FALLBACKS in favor of GenerationConfig dataclass defaults,
and streamline build_generation_config to merge request and HF config
in one place.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: lvhan028

lmdeploy/serve/anthropic/adapter.py

@@ -9,7 +9,7 @@
 import shortuuid
 
 from lmdeploy.messages import GenerationConfig
-from lmdeploy.serve.core.generation_config import build_generation_config, extract_request_sampling_values
+from lmdeploy.serve.core.generation_config import build_generation_config
 from lmdeploy.serve.openai.protocol import Tool, ToolChoice, ToolChoiceFuncName
 
 from .protocol import (
@@ -347,11 +347,10 @@ def to_generation_config(
     default_gen_config: dict | None = None,
 ) -> GenerationConfig:
     """Map Anthropic messages request to LMDeploy generation config."""
-    request_values = extract_request_sampling_values(request)
     return build_generation_config(
-        request_values,
+        request,
         default_gen_config or {},
-        max_tokens=request.max_tokens,
+        max_new_tokens=request.max_tokens,
         stop_words=request.stop_sequences,
         include_stop_str_in_output=request.include_stop_str_in_output or False,
         skip_special_tokens=True,

lmdeploy/serve/core/generation_config.py

@@ -4,30 +4,14 @@
 
 from __future__ import annotations
 
+import dataclasses
 from typing import Any
 
 from lmdeploy.messages import GenerationConfig
 from lmdeploy.utils import get_logger
 
 logger = get_logger('lmdeploy')
 
-PROTOCOL_FALLBACKS: dict[str, Any] = {
-    'temperature': 0.7,
-    'top_p': 1.0,
-    'top_k': 40,
-    'repetition_penalty': 1.0,
-    'min_p': 0.0,
-    'do_sample': True,
-}
-
-REQUEST_SAMPLING_FIELDS = (
-    'temperature',
-    'top_p',
-    'top_k',
-    'min_p',
-    'repetition_penalty',
-)
-
 
 def _load_hf_generation_config(path: str, trust_remote_code: bool) -> dict[str, Any]:
     from transformers import GenerationConfig
@@ -66,55 +50,45 @@ def resolve_default_gen_config(
 def merge_sampling_params(
     request_values: dict[str, Any],
     default_gen_config: dict[str, Any],
-    fallbacks: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
-    """Merge sampling params with request > default_gen_config > protocol
-    fallback priority."""
-    fallbacks = fallbacks or PROTOCOL_FALLBACKS
+    """Merge sampling params with request > default_gen_config priority."""
     merged: dict[str, Any] = {}
-    all_keys = set(fallbacks) | set(default_gen_config) | set(request_values)
-    for key in all_keys:
+    for key in set(default_gen_config) | set(request_values):
         if key in request_values:
             merged[key] = request_values[key]
-        elif key in default_gen_config:
+        else:
             merged[key] = default_gen_config[key]
-        elif key in fallbacks:
-            merged[key] = fallbacks[key]
     return merged
 
 
 def extract_request_sampling_values(request: Any) -> dict[str, Any]:
-    """Extract explicitly provided sampling fields from a request object."""
+    """Extract non-None GenerationConfig fields present on the request."""
     values: dict[str, Any] = {}
-    for field in REQUEST_SAMPLING_FIELDS:
-        if not hasattr(request, field):
+    for field in dataclasses.fields(GenerationConfig):
+        if not hasattr(request, field.name):
             continue
-        value = getattr(request, field)
+        value = getattr(request, field.name)
         if value is not None:
-            values[field] = value
+            values[field.name] = value
     return values
 
 
 def build_generation_config(
-    request_values: dict[str, Any],
+    request: Any,
     default_gen_config: dict[str, Any],
     *,
-    max_completion_tokens: int | None = None,
-    max_tokens: int | None = None,
-    fallbacks: dict[str, Any] | None = None,
+    max_new_tokens: int | None = None,
     **extra_kwargs: Any,
 ) -> GenerationConfig:
     """Build ``GenerationConfig`` from merged sampling defaults and request
     values."""
-    merged = merge_sampling_params(request_values, default_gen_config, fallbacks)
-    max_new_tokens = max_completion_tokens if max_completion_tokens is not None else max_tokens
+    request_values = extract_request_sampling_values(request)
+    merged = merge_sampling_params(request_values, default_gen_config)
+    merged.pop('max_new_tokens', None)
+    merged.pop('do_sample', None)
     return GenerationConfig(
         max_new_tokens=max_new_tokens,
-        do_sample=merged.get('do_sample', PROTOCOL_FALLBACKS['do_sample']),
-        top_k=merged['top_k'],
-        top_p=merged['top_p'],
-        temperature=merged['temperature'],
-        repetition_penalty=merged['repetition_penalty'],
-        min_p=merged['min_p'],
+        do_sample=True,
+        **merged,
         **extra_kwargs,
     )

lmdeploy/serve/openai/api_server.py

@@ -49,7 +49,6 @@
 from lmdeploy.serve.core import AsyncEngine, EngineHealthMonitor
 from lmdeploy.serve.core.generation_config import (
     build_generation_config,
-    extract_request_sampling_values,
     resolve_default_gen_config,
 )
 from lmdeploy.serve.openai.protocol import (
@@ -155,16 +154,13 @@ def get_engine_config(cls):
 
 def _build_serving_generation_config(request, **extra_kwargs) -> GenerationConfig:
     """Build ``GenerationConfig`` with server and request sampling merge."""
-    request_values = extract_request_sampling_values(request)
-    max_completion_tokens = getattr(request, 'max_completion_tokens', None)
-    max_tokens = getattr(request, 'max_tokens', None)
-    if max_completion_tokens is None and hasattr(request, 'max_output_tokens'):
-        max_completion_tokens = getattr(request, 'max_output_tokens', None)
+    max_new_tokens = getattr(request, 'max_completion_tokens', None)
+    if max_new_tokens is None:
+        max_new_tokens = getattr(request, 'max_tokens', None)
     return build_generation_config(
-        request_values,
+        request,
         VariableInterface.default_gen_config,
-        max_completion_tokens=max_completion_tokens,
-        max_tokens=max_tokens,
+        max_new_tokens=max_new_tokens,
         **extra_kwargs,
     )
 

lmdeploy/serve/openai/responses/request.py

@@ -9,7 +9,7 @@
 from fastapi.responses import JSONResponse
 
 from lmdeploy.messages import GenerationConfig
-from lmdeploy.serve.core.generation_config import build_generation_config, extract_request_sampling_values
+from lmdeploy.serve.core.generation_config import build_generation_config
 from lmdeploy.serve.openai.protocol import Tool, ToolChoice, ToolChoiceFuncName
 from lmdeploy.serve.openai.responses.protocol import ResponsesRequest
 from lmdeploy.utils import get_logger
@@ -270,11 +270,10 @@ def to_generation_config(
     default_gen_config: dict | None = None,
 ) -> GenerationConfig:
     stop_words = [request.stop] if isinstance(request.stop, str) else request.stop
-    request_values = extract_request_sampling_values(request)
     return build_generation_config(
-        request_values,
+        request,
         default_gen_config or {},
-        max_completion_tokens=request.max_output_tokens,
+        max_new_tokens=request.max_output_tokens,
         stop_words=stop_words,
         ignore_eos=request.ignore_eos,
         skip_special_tokens=request.skip_special_tokens,

lmdeploy/serve/openai/serving_chat_completion.py

@@ -1,11 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from typing import TYPE_CHECKING
 
-from lmdeploy.serve.core.generation_config import (
-    PROTOCOL_FALLBACKS,
-    extract_request_sampling_values,
-    merge_sampling_params,
-)
+from lmdeploy.serve.core.generation_config import build_generation_config
 
 from .protocol import ChatCompletionRequest
 
@@ -14,11 +10,12 @@
 
 
 def _effective_sampling(request: ChatCompletionRequest, server_context: 'VariableInterface') -> dict:
-    return merge_sampling_params(
-        extract_request_sampling_values(request),
-        server_context.default_gen_config,
-        PROTOCOL_FALLBACKS,
-    )
+    gen = build_generation_config(request, server_context.default_gen_config)
+    return {
+        'temperature': gen.temperature,
+        'top_p': gen.top_p,
+        'top_k': gen.top_k,
+    }
 
 
 def check_request(request: ChatCompletionRequest, server_context: 'VariableInterface') -> str:

lmdeploy/serve/openai/serving_completion.py

@@ -1,11 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from typing import TYPE_CHECKING
 
-from lmdeploy.serve.core.generation_config import (
-    PROTOCOL_FALLBACKS,
-    extract_request_sampling_values,
-    merge_sampling_params,
-)
+from lmdeploy.serve.core.generation_config import build_generation_config
 
 from .protocol import CompletionRequest
 
@@ -14,11 +10,12 @@
 
 
 def _effective_sampling(request: CompletionRequest, server_context: 'VariableInterface') -> dict:
-    return merge_sampling_params(
-        extract_request_sampling_values(request),
-        server_context.default_gen_config,
-        PROTOCOL_FALLBACKS,
-    )
+    gen = build_generation_config(request, server_context.default_gen_config)
+    return {
+        'temperature': gen.temperature,
+        'top_p': gen.top_p,
+        'top_k': gen.top_k,
+    }
 
 
 def check_request(request: CompletionRequest, server_context: 'VariableInterface') -> str:

tests/test_lmdeploy/serve/test_generation_config.py

@@ -1,31 +1,29 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from unittest.mock import patch
 
+from lmdeploy.messages import GenerationConfig
 from lmdeploy.serve.core.generation_config import (
-    PROTOCOL_FALLBACKS,
     build_generation_config,
     extract_request_sampling_values,
     merge_sampling_params,
     resolve_default_gen_config,
 )
 from lmdeploy.serve.openai.protocol import ChatCompletionRequest, CompletionRequest
 
+_DEFAULTS = GenerationConfig()
+
 
 def test_merge_sampling_params_priority():
     merged = merge_sampling_params(
         {'temperature': 0.2},
         {'temperature': 0.5, 'top_k': 10},
-        PROTOCOL_FALLBACKS,
     )
-    assert merged['temperature'] == 0.2
-    assert merged['top_k'] == 10
-    assert merged['top_p'] == PROTOCOL_FALLBACKS['top_p']
+    assert merged == {'temperature': 0.2, 'top_k': 10}
 
 
-def test_merge_sampling_params_uses_server_then_fallback():
-    merged = merge_sampling_params({}, {'temperature': 0.5}, PROTOCOL_FALLBACKS)
-    assert merged['temperature'] == 0.5
-    assert merged['top_k'] == PROTOCOL_FALLBACKS['top_k']
+def test_merge_sampling_params_uses_server_defaults():
+    merged = merge_sampling_params({}, {'temperature': 0.5})
+    assert merged == {'temperature': 0.5}
 
 
 def test_extract_request_sampling_values_only_non_null():
@@ -35,10 +33,11 @@ def test_extract_request_sampling_values_only_non_null():
 
 
 def test_build_generation_config_from_merged_values():
+    request = ChatCompletionRequest(model='test', messages='hi', temperature=0.2)
     gen_config = build_generation_config(
-        {'temperature': 0.2},
+        request,
         {'top_k': 5},
-        max_completion_tokens=32,
+        max_new_tokens=32,
     )
     assert gen_config.temperature == 0.2
     assert gen_config.top_k == 5
@@ -47,10 +46,18 @@ def test_build_generation_config_from_merged_values():
 
 
 def test_build_generation_config_max_new_tokens_defaults_to_none():
-    gen_config = build_generation_config({}, {})
+    request = CompletionRequest(model='test', prompt='hello')
+    gen_config = build_generation_config(request, {})
     assert gen_config.max_new_tokens is None
 
 
+def test_build_generation_config_uses_generation_config_defaults():
+    request = CompletionRequest(model='test', prompt='hello')
+    gen_config = build_generation_config(request, {})
+    assert gen_config.temperature == _DEFAULTS.temperature
+    assert gen_config.top_k == _DEFAULTS.top_k
+
+
 @patch('lmdeploy.serve.core.generation_config._load_hf_generation_config')
 def test_resolve_default_gen_config_auto(mock_load):
     mock_load.return_value = {
@@ -74,9 +81,6 @@ def test_resolve_default_gen_config_lmdeploy():
 
 def test_completion_request_sampling_merge():
     request = CompletionRequest(model='test', prompt='hello')
-    gen_config = build_generation_config(
-        extract_request_sampling_values(request),
-        {'temperature': 0.9},
-    )
+    gen_config = build_generation_config(request, {'temperature': 0.9})
     assert gen_config.temperature == 0.9
-    assert gen_config.top_k == PROTOCOL_FALLBACKS['top_k']
+    assert gen_config.top_k == _DEFAULTS.top_k