[EngineConfig] Add override options (#2550)

This PR introduces override options to the Python side EngineConfig
so that they'll be reflected in JIT model compilation.

Author: MasterJH5574

docs/deploy/cli.rst

@@ -87,5 +87,5 @@ MODEL                  The model folder after compiling with MLC-LLM build proce
                        with the device id set to 0 for default.
 --overrides            Model configuration override. Supports overriding
                        ``context_window_size``, ``prefill_chunk_size``, ``sliding_window_size``, ``attention_sink_size``,
-                       ``max_batch_size`` and ``tensor_parallel_shards``. The overrides could be explicitly
+                       and ``tensor_parallel_shards``. The overrides could be explicitly
                        specified via details knobs, e.g. --overrides ``context_window_size=1024;prefill_chunk_size=128``.

docs/deploy/python_engine.rst

@@ -94,12 +94,12 @@ for the complete chat completion interface.
   .. code:: python
 
     from mlc_llm import MLCEngine
-    from mlc_llm.serve.config import ModelConfigOverride
+    from mlc_llm.serve.config import EngineConfig
 
     model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC"
     engine = MLCEngine(
         model,
-        model_config_overrides=ModelConfigOverride(tensor_parallel_shards=2),
+        engine_config=EngineConfig(tensor_parallel_shards=2),
     )
 
 
@@ -196,12 +196,12 @@ for the complete chat completion interface.
   .. code:: python
 
     from mlc_llm import AsyncMLCEngine
-    from mlc_llm.serve.config import ModelConfigOverride
+    from mlc_llm.serve.config import EngineConfig
 
     model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC"
     engine = AsyncMLCEngine(
         model,
-        model_config_overrides=ModelConfigOverride(tensor_parallel_shards=2),
+        engine_config=EngineConfig(tensor_parallel_shards=2),
     )
 
 

docs/get_started/introduction.rst

@@ -153,12 +153,12 @@ If you would like to do concurrent asynchronous generation, you can use :class:`
   .. code:: python
 
     from mlc_llm import MLCEngine
-    from mlc_llm.serve.config import ModelConfigOverride
+    from mlc_llm.serve.config import EngineConfig
 
     model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC"
     engine = MLCEngine(
         model,
-        model_config_overrides=ModelConfigOverride(tensor_parallel_shards=2),
+        engine_config=EngineConfig(tensor_parallel_shards=2),
     )
 
 

python/mlc_llm/cli/calibrate.py

@@ -4,7 +4,7 @@
 from mlc_llm.interface.help import HELP
 from mlc_llm.support.argparse import ArgumentParser
 
-from .serve import EngineAndModelConfigOverride
+from .serve import EngineConfigOverride
 
 
 def main(argv):
@@ -51,7 +51,7 @@ def main(argv):
     )
     parser.add_argument(
         "--overrides",
-        type=EngineAndModelConfigOverride.from_str,
+        type=EngineConfigOverride.from_str,
         default="",
         help=HELP["overrides_serve"],
     )

python/mlc_llm/cli/chat.py

@@ -1,8 +1,7 @@
 """Command line entrypoint of chat."""
 
-from mlc_llm.interface.chat import chat
+from mlc_llm.interface.chat import ModelConfigOverride, chat
 from mlc_llm.interface.help import HELP
-from mlc_llm.serve.config import ModelConfigOverride
 from mlc_llm.support.argparse import ArgumentParser
 
 

python/mlc_llm/cli/compile.py

@@ -25,7 +25,7 @@
 
 
 def main(argv):
-    """Parse command line argumennts and call `mlc_llm.compiler.compile`."""
+    """Parse command line arguments and call `mlc_llm.compiler.compile`."""
 
     def _parse_output(path: Union[str, Path]) -> Path:
         path = Path(path)

python/mlc_llm/cli/serve.py

@@ -7,13 +7,12 @@
 
 from mlc_llm.interface.help import HELP
 from mlc_llm.interface.serve import serve
-from mlc_llm.serve.config import ModelConfigOverride
 from mlc_llm.support import argparse
 from mlc_llm.support.argparse import ArgumentParser
 
 
 @dataclasses.dataclass
-class EngineAndModelConfigOverride:  # pylint: disable=too-many-instance-attributes
+class EngineConfigOverride:  # pylint: disable=too-many-instance-attributes
     """Arguments for overriding engine config."""
 
     # Overrides for EngineConfig (runtime)
@@ -24,8 +23,6 @@ class EngineAndModelConfigOverride:  # pylint: disable=too-many-instance-attribu
     gpu_memory_utilization: Optional[float] = None
     spec_draft_length: Optional[int] = None
     prefix_cache_max_num_recycling_seqs: Optional[int] = None
-
-    # Overrides for model config (compile time)
     context_window_size: Optional[int] = None
     sliding_window_size: Optional[int] = None
     attention_sink_size: Optional[int] = None
@@ -51,7 +48,7 @@ def __repr__(self) -> str:
         return out.getvalue().rstrip()
 
     @staticmethod
-    def from_str(source: str) -> "EngineAndModelConfigOverride":
+    def from_str(source: str) -> "EngineConfigOverride":
         """Parse engine config override values from a string."""
         parser = argparse.ArgumentParser(description="Engine config override values")
 
@@ -67,7 +64,7 @@ def from_str(source: str) -> "EngineAndModelConfigOverride":
         parser.add_argument("--attention_sink_size", type=int, default=None)
         parser.add_argument("--tensor_parallel_shards", type=int, default=None)
         results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
-        return EngineAndModelConfigOverride(
+        return EngineConfigOverride(
             max_num_sequence=results.max_num_sequence,
             max_total_seq_length=results.max_total_seq_length,
             prefill_chunk_size=results.prefill_chunk_size,
@@ -81,17 +78,6 @@ def from_str(source: str) -> "EngineAndModelConfigOverride":
             tensor_parallel_shards=results.tensor_parallel_shards,
         )
 
-    def to_model_config_overrides(self) -> ModelConfigOverride:
-        """Extract the model config overrides."""
-        return ModelConfigOverride(
-            context_window_size=self.context_window_size,
-            sliding_window_size=self.sliding_window_size,
-            prefill_chunk_size=self.prefill_chunk_size,
-            attention_sink_size=self.attention_sink_size,
-            max_batch_size=self.max_num_sequence,
-            tensor_parallel_shards=self.tensor_parallel_shards,
-        )
-
 
 def main(argv):
     """Parse command line arguments and call `mlc_llm.interface.serve`."""
@@ -145,7 +131,7 @@ def main(argv):
     )
     parser.add_argument(
         "--overrides",
-        type=EngineAndModelConfigOverride.from_str,
+        type=EngineConfigOverride.from_str,
         default="",
         help=HELP["overrides_serve"],
     )
@@ -199,16 +185,19 @@ def main(argv):
         mode=parsed.mode,
         enable_debug=parsed.enable_debug,
         additional_models=additional_models,
+        tensor_parallel_shards=parsed.overrides.tensor_parallel_shards,
         speculative_mode=parsed.speculative_mode,
         prefix_cache_mode=parsed.prefix_cache_mode,
         max_num_sequence=parsed.overrides.max_num_sequence,
         max_total_sequence_length=parsed.overrides.max_total_seq_length,
+        max_single_sequence_length=parsed.overrides.context_window_size,
         prefill_chunk_size=parsed.overrides.prefill_chunk_size,
+        sliding_window_size=parsed.overrides.sliding_window_size,
+        attention_sink_size=parsed.overrides.attention_sink_size,
         max_history_size=parsed.overrides.max_history_size,
         gpu_memory_utilization=parsed.overrides.gpu_memory_utilization,
         spec_draft_length=parsed.overrides.spec_draft_length,
         prefix_cache_max_num_recycling_seqs=parsed.overrides.prefix_cache_max_num_recycling_seqs,
-        model_config_overrides=parsed.overrides.to_model_config_overrides(),
         enable_tracing=parsed.enable_tracing,
         host=parsed.host,
         port=parsed.port,

python/mlc_llm/interface/chat.py

@@ -8,7 +8,7 @@
 
 from mlc_llm.json_ffi import JSONFFIEngine
 from mlc_llm.protocol import openai_api_protocol
-from mlc_llm.serve.config import EngineConfig, ModelConfigOverride
+from mlc_llm.serve.config import EngineConfig
 from mlc_llm.serve.engine import MLCEngine
 from mlc_llm.serve.engine_base import _query_engine_metrics
 from mlc_llm.support import argparse
@@ -79,6 +79,36 @@ def from_str(source: str) -> "ChatCompletionOverride":
         )
 
 
+@dataclasses.dataclass
+class ModelConfigOverride(ConfigOverrideBase):  # pylint: disable=too-many-instance-attributes
+    """Flags for overriding model config."""
+
+    context_window_size: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    prefill_chunk_size: Optional[int] = None
+    attention_sink_size: Optional[int] = None
+    tensor_parallel_shards: Optional[int] = None
+
+    @staticmethod
+    def from_str(source: str) -> "ModelConfigOverride":
+        """Parse model config override values from a string."""
+        parser = argparse.ArgumentParser(description="model config override values")
+        parser.add_argument("--tensor_parallel_shards", type=int, default=None)
+        parser.add_argument("--context_window_size", type=int, default=None)
+        parser.add_argument("--sliding_window_size", type=int, default=None)
+        parser.add_argument("--prefill_chunk_size", type=int, default=None)
+        parser.add_argument("--attention_sink_size", type=int, default=None)
+
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return ModelConfigOverride(
+            tensor_parallel_shards=results.tensor_parallel_shards,
+            context_window_size=results.context_window_size,
+            sliding_window_size=results.sliding_window_size,
+            prefill_chunk_size=results.prefill_chunk_size,
+            attention_sink_size=results.attention_sink_size,
+        )
+
+
 class ChatState:
     """Simple helper class to manage chat state.
 
@@ -255,8 +285,11 @@ def chat(
             model_lib=model_lib,
             mode="interactive",
             engine_config=EngineConfig(
+                max_single_sequence_length=overrides.context_window_size,
                 prefill_chunk_size=overrides.prefill_chunk_size,
+                sliding_window_size=overrides.sliding_window_size,
+                attention_sink_size=overrides.attention_sink_size,
+                tensor_parallel_shards=overrides.tensor_parallel_shards,
             ),
-            model_config_overrides=overrides,
         )
     ).chat()

python/mlc_llm/interface/help.py

@@ -128,7 +128,7 @@
     "modelconfig_overrides": """
 Model configuration override. Supports overriding,
 `context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`,
-`max_batch_size` and `tensor_parallel_shards`. The overrides could be explicitly
+`max_num_sequence` and `tensor_parallel_shards`. The overrides could be explicitly
 specified via details knobs, e.g. --overrides "context_window_size=1024;prefill_chunk_size=128".
 """.strip(),
     "debug_dump": """

python/mlc_llm/interface/serve.py

@@ -8,7 +8,6 @@
 
 from mlc_llm.protocol import error_protocol
 from mlc_llm.serve import engine
-from mlc_llm.serve.config import ModelConfigOverride
 from mlc_llm.serve.entrypoints import (
     debug_entrypoints,
     metrics_entrypoints,
@@ -27,16 +26,19 @@ def serve(
     mode: Literal["local", "interactive", "server"],
     enable_debug: bool,
     additional_models: List[Union[str, Tuple[str, str]]],
+    tensor_parallel_shards: Optional[int],
     max_num_sequence: Optional[int],
     max_total_sequence_length: Optional[int],
+    max_single_sequence_length: Optional[int],
     prefill_chunk_size: Optional[int],
+    sliding_window_size: Optional[int],
+    attention_sink_size: Optional[int],
     max_history_size: Optional[int],
     gpu_memory_utilization: Optional[float],
     speculative_mode: Literal["disable", "small_draft", "eagle", "medusa"],
     spec_draft_length: Optional[int],
     prefix_cache_mode: Literal["disable", "radix"],
     prefix_cache_max_num_recycling_seqs: Optional[int],
-    model_config_overrides: Optional[ModelConfigOverride],
     enable_tracing: bool,
     host: str,
     port: int,
@@ -54,17 +56,20 @@ def serve(
         mode=mode,
         engine_config=engine.EngineConfig(
             additional_models=additional_models,
+            tensor_parallel_shards=tensor_parallel_shards,
             max_num_sequence=max_num_sequence,
             max_total_sequence_length=max_total_sequence_length,
+            max_single_sequence_length=max_single_sequence_length,
             prefill_chunk_size=prefill_chunk_size,
+            sliding_window_size=sliding_window_size,
+            attention_sink_size=attention_sink_size,
             max_history_size=max_history_size,
             gpu_memory_utilization=gpu_memory_utilization,
             speculative_mode=speculative_mode,
             spec_draft_length=spec_draft_length,
             prefix_cache_mode=prefix_cache_mode,
             prefix_cache_max_num_recycling_seqs=prefix_cache_max_num_recycling_seqs,
         ),
-        model_config_overrides=model_config_overrides,
         enable_tracing=enable_tracing,
     )