chore: import upstream snapshot with attribution

python/mlc_llm/interface/calibrate.py

@@ -0,0 +1,172 @@
+"""Python entrypoint for calibration."""
+
+import asyncio
+import json
+import random
+from collections.abc import Mapping
+from typing import List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+import tqdm.asyncio
+import tvm
+from tvm.contrib import tvmjs
+
+from mlc_llm.serve.engine import AsyncMLCEngine, EngineConfig
+from mlc_llm.tokenizers import Tokenizer
+
+
+class CalibrationObserver:
+    """A singleton class to observe the calibration parameters.""" ""
+
+    instance: "CalibrationObserver" = None
+
+    params: Mapping[str, tvm.runtime.Tensor] = {}
+
+    @staticmethod
+    def get():
+        """Get the singleton instance of the class.""" ""
+        if CalibrationObserver.instance is None:
+            CalibrationObserver.instance = CalibrationObserver()
+        return CalibrationObserver.instance
+
+    @tvm.register_global_func("mlc_llm.calibration_observer")
+    @staticmethod
+    def callback(
+        name: str,
+        mode: str,
+        value: "tvm.runtime.Tensor",
+        out_value: "tvm.runtime.Tensor",
+    ):
+        """The callback function to update the saved calibration parameters."""
+        instance = CalibrationObserver.get()
+        if mode == "max":
+            reducer = np.maximum
+        else:
+            raise NotImplementedError(f"Unsupported calibration mode: {mode}")
+        if name in instance.params:
+            instance.params[name] = reducer(instance.params[name], value.numpy())
+        else:
+            instance.params[name] = value.numpy()
+        out_value.copyfrom(instance.params[name])
+
+    def save_params(self, output: str):
+        """Save the calibration parameters to the given output directory."""
+        tvmjs.dump_tensor_cache(
+            self.params,
+            output,
+            encode_format="f32-to-bf16",
+            meta_data=None,
+            show_progress=False,
+            update_if_exists=True,
+        )
+
+
+def sample_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: Tokenizer,
+) -> List[Tuple[str, int, int]]:  # noqa: UP006
+    """Sample the requests from the given dataset."""
+    # Load the dataset.
+    with open(dataset_path, encoding="utf-8") as f:
+        dataset = json.load(f)
+
+    # Filter out the conversations with less than 2 turns.
+    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
+    # Only keep the first two turns of each conversation.
+    dataset = [
+        (data["conversations"][0]["value"], data["conversations"][1]["value"]) for data in dataset
+    ]
+    prompts = [prompt for prompt, _ in dataset]
+    prompt_token_ids = tokenizer.encode_batch(prompts)
+    completions = [completion for _, completion in dataset]
+    completion_token_ids = tokenizer.encode_batch(completions)
+    tokenized_dataset: List[Tuple[str, List[int], int]] = []  # noqa: UP006
+    for i in range(len(dataset)):
+        output_len = len(completion_token_ids[i])
+        tokenized_dataset.append((prompts[i], prompt_token_ids[i], output_len))
+
+    # Filter out too long sequences.
+    filtered_dataset: List[Tuple[str, int, int]] = []  # noqa: UP006
+    for prompt, token_ids, output_len in tokenized_dataset:
+        prompt_len = len(token_ids)
+        if prompt_len < 4 or output_len < 4:
+            # Prune too short sequences.
+            continue
+        if prompt_len > 1024 or prompt_len + output_len > 2048:
+            # Prune too long sequences.
+            continue
+        filtered_dataset.append((prompt, prompt_len, output_len))
+
+    # Sample the requests.
+    sampled_requests = random.sample(filtered_dataset, num_requests)
+    return sampled_requests
+
+
+async def send_calibration_requests(
+    async_engine: AsyncMLCEngine,
+    sampled_requests: List[Tuple[str, int, int]],  # noqa: UP006
+    max_concurrent_requests: int,
+) -> None:
+    """Send the calibration requests to the engine."""
+    tasks = []
+
+    semaphore = asyncio.Semaphore(max_concurrent_requests)
+
+    async def generate_task(request_idx):
+        async with semaphore:
+            prompt, _, output_len = sampled_requests[request_idx]
+            await async_engine.chat.completions.create(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=output_len,
+                request_id=str(request_idx),
+            )
+
+    for i in range(len(sampled_requests)):
+        task = asyncio.create_task(generate_task(i))
+        tasks.append(task)
+    await tqdm.asyncio.tqdm.gather(*tasks)
+
+
+def calibrate(
+    model: str,
+    device: str,
+    model_lib: Optional[str],
+    dataset: str,
+    output: str,
+    num_calibration_samples: int,
+    *,
+    seed: int,
+    max_num_sequence: Optional[int] = None,
+    max_total_sequence_length: Optional[int] = None,
+    prefill_chunk_size: Optional[int] = None,
+    max_history_size: Optional[int] = None,
+    gpu_memory_utilization: Optional[float] = None,
+) -> None:
+    """Calibrate the quantized model using the given dataset."""
+    random.seed(seed)
+    async_engine = AsyncMLCEngine(
+        model=model,
+        device=device,
+        model_lib=model_lib,
+        mode="server",
+        engine_config=EngineConfig(
+            max_num_sequence=max_history_size,
+            max_total_sequence_length=max_total_sequence_length,
+            prefill_chunk_size=prefill_chunk_size,
+            max_history_size=max_history_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+        ),
+    )
+    sampled_requests = sample_requests(dataset, num_calibration_samples, async_engine.tokenizer)
+    asyncio.run(
+        send_calibration_requests(
+            async_engine,
+            sampled_requests,
+            max_concurrent_requests=max_num_sequence or 32,
+        )
+    )
+    async_engine.terminate()
+
+    calibrator = CalibrationObserver.get()
+    calibrator.save_params(output)

python/mlc_llm/interface/chat.py

@@ -0,0 +1,311 @@
+"""Python entrypoint of chat."""
+
+import dataclasses
+from typing import Any, Dict, List, Optional, Union  # noqa: UP035
+
+from prompt_toolkit import prompt as get_prompt
+from prompt_toolkit.key_binding import KeyBindings
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.protocol import openai_api_protocol
+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
+from mlc_llm.support.config import ConfigOverrideBase
+
+
+def _print_help_str():
+    help_str = """You can use the following special commands:
+  /help               print the special commands
+  /exit               quit the cli
+  /stats              print out stats of last request (token/sec)
+  /metrics            print out full engine metrics
+  /reset              restart a fresh chat
+  /set [overrides]    override settings in the generation config. For example,
+                      `/set temperature=0.5;top_p=0.8;seed=23;max_tokens=100;stop=str1,str2`
+                      Note: Separate stop words in the `stop` option with commas (,).
+  Multi-line input: Use escape+enter to start a new line.
+"""
+    print(help_str)
+
+
+def _set_up_key_bindings():
+    kb = KeyBindings()
+
+    @kb.add("escape", "enter")
+    def _(event):
+        event.current_buffer.insert_text("\n")
+
+    @kb.add("enter")
+    def _(event):
+        event.current_buffer.validate_and_handle()
+
+    return kb
+
+
+@dataclasses.dataclass
+class ChatCompletionOverride(ConfigOverrideBase):
+    """Flags for overriding chat completions."""
+
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    frequency_penalty: Optional[float] = None
+    presence_penalty: Optional[float] = None
+    max_tokens: Optional[int] = None
+    seed: Optional[int] = None
+    stop: Optional[Union[str, List[str]]] = None  # noqa: UP006
+
+    @staticmethod
+    def from_str(source: str) -> "ChatCompletionOverride":
+        """Parse model config override values from a string."""
+        parser = argparse.ArgumentParser(description="chat completion override values")
+        parser.add_argument("--temperature", type=float, default=None)
+        parser.add_argument("--top_p", type=float, default=None)
+        parser.add_argument("--frequency_penalty", type=float, default=None)
+        parser.add_argument("--presence_penalty", type=float, default=None)
+        parser.add_argument("--max_tokens", type=int, default=None)
+        parser.add_argument("--seed", type=int, default=None)
+        parser.add_argument("--stop", type=str, default=None)
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return ChatCompletionOverride(
+            temperature=results.temperature,
+            top_p=results.top_p,
+            frequency_penalty=results.frequency_penalty,
+            presence_penalty=results.presence_penalty,
+            max_tokens=results.max_tokens,
+            seed=results.seed,
+            stop=results.stop.split(",") if results.stop is not None else None,
+        )
+
+
+@dataclasses.dataclass
+class ModelConfigOverride(ConfigOverrideBase):
+    """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
+    pipeline_parallel_stages: Optional[int] = None
+    opt: Optional[str] = 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("--pipeline_parallel_stages", type=int, default=None)
+        parser.add_argument("--opt", type=str, 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,
+            pipeline_parallel_stages=results.pipeline_parallel_stages,
+            opt=results.opt,
+            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.
+
+    Chat state wraps around a  engine instance
+    and exposes the minimum set of tools to perform
+    interactive chat. It provides support for mlc_llm chat.
+    It also can be used to do interactive debugging
+    with different engine instance.
+
+    Examples
+    --------
+    .. code:: python
+
+        from openai import OpenAI
+        from mlc_llm import MLCEngine
+        from mlc_llm.serve import PopenServer
+        from mlc_llm.interface.chat import ChatState
+
+        def chat_with_engine(model):
+            # hookup with MLCEngine
+            ChatState(MLCEngine(model)).chat()
+
+        def chat_with_server(model):
+            # hookup with AsyncMLCEngine backed api server
+            with PopenServer(model) as server:
+                ChatState(
+                    OpenAI(base_url=server.openai_v1_base_url, api_key="None")
+                ).chat()
+    """
+
+    history: List[Dict[str, Any]]  # noqa: UP006
+    history_begin: int
+    # kwargs passed to completions
+    overrides: ChatCompletionOverride
+    # Underlying engine
+    engine: Union[JSONFFIEngine, MLCEngine]
+    last_finished_request_usage: Optional[openai_api_protocol.CompletionUsage]
+
+    def __init__(self, engine: Union[JSONFFIEngine, MLCEngine]):
+        self.engine = engine
+        self.history = []
+        self.history_window_begin = 0
+        self.overrides = ChatCompletionOverride()
+        # model is mainly used for compact reasons
+        self.model = "chat_model"
+        self.last_finished_request_usage = None
+
+    def slide_history(self):
+        """Slide history to fit into context window"""
+        history_window_size = len(self.history) - self.history_window_begin
+        assert history_window_size % 2 == 0
+        self.history_window_begin += ((history_window_size + 3) // 4) * 2
+
+    def process_system_prompts(self):
+        """Process system prompts"""
+        # TODO(mlc-team): possibly leverage debug option
+        # pass a simple prompt to warm up
+        for _ in self.engine.chat.completions.create(
+            messages=[{"role": "user", "content": ""}],
+            max_tokens=1,
+            model=self.model,
+            stream=True,
+        ):
+            pass
+
+    def generate(self, prompt: str):
+        """Run one generation with the prompt.
+
+        Parameters
+        ----------
+        prompt: str
+            The input prompt
+        """
+        self.history.append({"role": "user", "content": prompt})
+        output_text = ""
+        finish_reason_length = False
+        messages = self.history[self.history_window_begin :]
+
+        for response in self.engine.chat.completions.create(
+            messages=messages,
+            model=self.model,
+            stream=True,
+            stream_options={"include_usage": True},
+            **dataclasses.asdict(self.overrides),
+        ):
+            if response.usage is not None:
+                self.last_finished_request_usage = response.usage
+                continue
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                if isinstance(choice.delta.content, str):
+                    output_text += choice.delta.content
+                    print(choice.delta.content, end="", flush=True)
+                if choice.finish_reason == "length":
+                    finish_reason_length = True
+        if finish_reason_length:
+            print(" [output truncated due to context length limit...]")
+        # print additional \n when generation ends
+        print()
+        # record the history
+        self.history.append({"role": "assistant", "content": output_text})
+        if finish_reason_length:
+            self.slide_history()
+
+    def stats(self):
+        """Print statistics of the prefill and decode speed."""
+
+        def get_stats_text():
+            """Get text"""
+            if self.last_finished_request_usage is None:
+                return "N/A"
+            last_finished_request = self.last_finished_request_usage.extra
+            if last_finished_request is None:
+                return "N/A"
+            prefill_speed = last_finished_request.get("prefill_tokens_per_s", None)
+            decode_speed = last_finished_request.get("decode_tokens_per_s", None)
+            prefill_speed = f"{prefill_speed:.1f}" if prefill_speed is not None else "N/A"
+            decode_speed = f"{decode_speed:.1f}" if decode_speed is not None else "N/A"
+            return f"prefill: {prefill_speed} tok/s, decode: {decode_speed} tok/s"
+
+        print(get_stats_text(), flush=True)
+
+    def metrics(self):
+        """Print metrics as prometheus text"""
+        print(_query_engine_metrics(self.engine).prometheus_text(), flush=True)
+
+    def reset(self):
+        """Reset the chat history"""
+        self.history = []
+        self.history_window_begin = 0
+
+    def chat(self):
+        """Start an interactive chat session."""
+        _print_help_str()
+
+        self.process_system_prompts()
+        # Multi-line input support: set escape+enter as start a new line
+        kb = _set_up_key_bindings()
+
+        while True:
+            try:
+                prompt = get_prompt(
+                    ">>> ",
+                    key_bindings=kb,
+                    multiline=True,
+                )
+            except (KeyboardInterrupt, EOFError):
+                break
+            if prompt[:4] == "/set":
+                overrides = ChatCompletionOverride.from_str(prompt.split()[1])
+                for key, value in dataclasses.asdict(overrides).items():
+                    if value is not None:
+                        setattr(self.overrides, key, value)
+            elif prompt[:6] == "/stats":
+                self.stats()
+            elif prompt[:8] == "/metrics":
+                self.metrics()
+            elif prompt[:6] == "/reset":
+                self.reset()
+            elif prompt[:5] == "/exit":
+                break
+            elif prompt[:5] == "/help":
+                _print_help_str()
+            else:
+                self.generate(prompt)
+
+
+def chat(
+    model: str,
+    device: str,
+    model_lib: Optional[str],
+    overrides: ModelConfigOverride,
+):
+    """Chat cli entry"""
+    # By default we use JSONFFIEngine
+    engine = JSONFFIEngine(
+        model,
+        device,
+        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,
+            pipeline_parallel_stages=overrides.pipeline_parallel_stages,
+            opt=overrides.opt,
+        ),
+    )
+    try:
+        ChatState(engine).chat()
+    finally:
+        engine.terminate()

python/mlc_llm/interface/compile.py

@@ -0,0 +1,265 @@
+"""Python entrypoint of compilation."""
+
+import dataclasses
+from io import StringIO
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Tuple  # noqa: UP035
+
+from tvm import IRModule, relax, tirx
+from tvm.ir.transform import Pass, PassContext
+from tvm.relax.frontend import nn
+from tvm.target import Target
+
+from mlc_llm import compiler_pass as _  # noqa: F401
+from mlc_llm import op as op_ext
+from mlc_llm.cli.model_metadata import _report_memory_usage
+from mlc_llm.model import Model
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import logging
+from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.style import bold
+
+from .compiler_flags import ModelConfigOverride, OptimizationFlags
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class CompileArgs:
+    """Arguments to MLC LLM's compiler."""
+
+    config: Path
+    quantization: Quantization
+    model: Model
+    target: Target
+    opt: OptimizationFlags
+    build_func: Callable[[IRModule, "CompileArgs", Pass], None]
+    system_lib_prefix: str
+    output: Path
+    overrides: ModelConfigOverride
+    debug_dump: Optional[Path]
+
+    def __post_init__(self) -> None:
+        self.opt.update(self.target, self.quantization)
+
+    def display(self) -> None:
+        """Display the arguments to stdout."""
+        out = StringIO()
+        print(f"{bold('Compiling with arguments:')}", file=out)
+        print(f"  {bold('--config'):<25} {self.config}", file=out)
+        print(f"  {bold('--quantization'):<25} {self.quantization}", file=out)
+        print(f"  {bold('--model-type'):<25} {self.model.name}", file=out)
+        print(f"  {bold('--target'):<25} {self.target.export()}", file=out)
+        print(f"  {bold('--opt'):<25} {self.opt}", file=out)
+        print(f'  {bold("--system-lib-prefix"):<25} "{self.system_lib_prefix}"', file=out)
+        print(f"  {bold('--output'):<25} {self.output}", file=out)
+        print(f"  {bold('--overrides'):<25} {self.overrides}", file=out)
+        # As it's debug only, no need to display
+        # print(f"  {bold('--debug-dump'):<25} {self.debug_dump}", file=out)
+        print(out.getvalue().rstrip())
+
+
+def _apply_preproc_to_params_and_check_pipeline(
+    named_params: List[Tuple[str, nn.Parameter]],  # noqa: UP006
+    model_config,
+) -> Dict[str, tirx.PrimFunc]:  # noqa: UP006
+    extra_tirs: Dict[str, tirx.PrimFunc] = {}  # noqa: UP006
+    for name, param in named_params:
+        preprocs = param.attrs.get("preprocs", [])
+        shard_strategy = param.attrs.get("shard_strategy", None)
+        if shard_strategy is not None and model_config.tensor_parallel_shards > 1:
+            preprocs.append(
+                shard_strategy.gen_shard_info(
+                    shards=model_config.tensor_parallel_shards,
+                    weight=param,
+                )
+            )
+            if shard_strategy.name not in extra_tirs:
+                extra_tirs[shard_strategy.name] = shard_strategy.gen_tir(
+                    shards=model_config.tensor_parallel_shards,
+                    weight=param,
+                )
+        param.attrs["preprocs"] = preprocs
+
+        pipeline_parallel_stages = getattr(model_config, "pipeline_parallel_stages", 1)
+        if pipeline_parallel_stages != 1:
+            assert "pipeline_stages" in param.attrs, (
+                f'The pipeline stage is undefined for parameter "{name}" when the number '
+                f"of pipeline parallel stages is {pipeline_parallel_stages}"
+            )
+        param.attrs["pipeline_stages"] = (
+            [0]
+            if "pipeline_stages" not in param.attrs
+            else list(set(param.attrs["pipeline_stages"]))
+        )
+    return extra_tirs
+
+
+def _infer_kv_state_kind(model_type) -> str:
+    if "rwkv" in model_type:
+        return "rnn_state"
+    if "medusa" in model_type:
+        return "none"
+    if "qwen3_5" in model_type:
+        return "hybrid"
+    return "kv_cache"
+
+
+def _compile(args: CompileArgs, model_config: ConfigBase):
+    def _get_variable_bounds(model_config) -> Dict[str, int]:  # noqa: UP006
+        if hasattr(model_config, "sliding_window_size"):
+            return {
+                "rolling_cache_len": model_config.sliding_window_size,
+                "kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size,
+                "seq_len": model_config.prefill_chunk_size,
+                "batch_size": getattr(model_config, "max_batch_size", 1),
+            }
+        return {
+            "total_seq_len": model_config.context_window_size,
+            "seq_len": model_config.prefill_chunk_size,
+            "batch_size": getattr(model_config, "max_batch_size", 1),
+        }
+
+    def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:  # noqa: UP006
+        return {
+            "name": name,
+            # Record dynamic shape as -1 (e.g. vocab_size)
+            "shape": [s if isinstance(s, int) else s.name for s in param.shape],
+            "dtype": str(param.dtype),
+            "preprocs": param.attrs["preprocs"],
+            "pipeline_stages": param.attrs.get("pipeline_stages", [0]),
+        }
+
+    logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
+    _kwargs = getattr(model_config, "kwargs", {})
+    model_config = args.overrides.apply(model_config)
+    with args.target:
+        op_ext.enable(
+            target=args.target,
+            flashinfer=args.opt.flashinfer,
+            faster_transformer=args.opt.faster_transformer,
+            cutlass=args.opt.cutlass,
+        )
+        # Step 1. Create the quantized model
+        logger.info("Creating model from: %s", model_config)
+        if (
+            args.quantization.kind == "ft-quant"
+            and hasattr(model_config, "tensor_parallel_shards")
+            and model_config.tensor_parallel_shards > 1
+        ):
+            raise NotImplementedError
+        if (
+            hasattr(args.quantization, "linear_weight_layout")
+            and args.quantization.linear_weight_layout == "KN"
+            and hasattr(model_config, "tensor_parallel_shards")
+            and model_config.tensor_parallel_shards > 1
+        ):
+            raise NotImplementedError(
+                "KN layout (q3f16_0 and q4f16_0) is not supported for tensor parallelism"
+            )
+        model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization)
+        # Step 2. Exporting the model to TVM
+        logger.info("Exporting the model to TVM compiler")
+        mod, named_params, ext_mods = model.export_tvm(
+            spec=model.get_default_spec(),
+            allow_extern=True,
+        )
+        # Step 3. Running relax compilation pipeline
+        logger.info("Running optimizations using TVM")
+        additional_tirs = _apply_preproc_to_params_and_check_pipeline(named_params, model_config)
+        variable_bounds = _get_variable_bounds(model_config)
+        cuda_graph_symbolic_capture_hints = {
+            "batch_decode": ["batch_size"],
+            "batch_decode_to_last_hidden_states": ["batch_size"],
+            "batch_verify": ["batch_size", "seq_len"],
+            "batch_verify_to_last_hidden_states": ["batch_size", "seq_len"],
+        }
+        avs = _kwargs.get("active_vocab_size", None)
+        if avs is not None and avs <= 0:
+            avs = None
+        metadata = {
+            "model_type": args.model.name,
+            "quantization": args.quantization.name,
+            "context_window_size": getattr(model_config, "context_window_size", -1),
+            "sliding_window_size": getattr(model_config, "sliding_window_size", -1),
+            "attention_sink_size": getattr(model_config, "attention_sink_size", -1),
+            "prefill_chunk_size": model_config.prefill_chunk_size,
+            "tensor_parallel_shards": model_config.tensor_parallel_shards,
+            "pipeline_parallel_stages": getattr(model_config, "pipeline_parallel_stages", 1),
+            "disaggregation": getattr(model_config, "disaggregation", False),
+            "kv_state_kind": _infer_kv_state_kind(args.model.name),
+            "max_batch_size": getattr(model_config, "max_batch_size", 1),
+            "active_vocab_size": avs,
+            "model_task": args.model.model_task,
+        }
+        if args.model.embedding_metadata:
+            metadata["embedding_metadata"] = dataclasses.asdict(args.model.embedding_metadata)
+        logger.info("Registering metadata: %s", metadata)
+        metadata["params"] = [_get_param_metadata(name, param) for name, param in named_params]
+        pass_config = {"relax.backend.use_cuda_graph": args.opt.cudagraph}
+        # TODO: Remove this workaround when the TVM CSE regression is fixed.
+        # Temporary workaround for TVM CSE regression that can produce
+        # dangling `cse_v*` vars during host codegen.
+        pass_config["tirx.disable_cse_tir"] = True
+
+        with PassContext(config=pass_config):
+            args.build_func(
+                mod,
+                args,
+                pipeline=relax.get_pipeline(
+                    "mlc_llm",
+                    target=args.target,
+                    flashinfer=args.opt.flashinfer,
+                    cublas_gemm=args.opt.cublas_gemm,
+                    faster_transformer=args.opt.faster_transformer,
+                    allreduce_strategy=args.opt.ipc_allreduce_strategy,
+                    variable_bounds=variable_bounds,
+                    cuda_graph_symbolic_capture_hints=cuda_graph_symbolic_capture_hints,
+                    additional_tirs=additional_tirs,
+                    ext_mods=ext_mods,
+                    metadata=metadata,
+                    debug_dump=args.debug_dump,
+                ),
+            )
+        _report_memory_usage(metadata=metadata, config=model_config)
+    logger.info("Generated: %s", bold(str(args.output)))
+
+
+def compile(
+    config: Dict[str, Any],  # noqa: UP006
+    quantization: Quantization,
+    model_type: Model,
+    target: Target,
+    opt: OptimizationFlags,
+    build_func: Callable[[IRModule, CompileArgs, Pass], None],
+    system_lib_prefix: str,
+    output: Path,
+    overrides: ModelConfigOverride,
+    debug_dump: Optional[Path] = None,
+):
+    """Compile a model given its configuration and quantization format to a specific target."""
+    avs = None
+    if "active_vocab_size" in config:
+        avs = config.pop("active_vocab_size")
+        logger.info("Active vocab size from input config: %s", str(avs))
+    if "model_config" in config:
+        model_config = config.pop("model_config")
+        model_config.update(config)
+        model_config = model_type.config.from_dict(model_config)
+    else:
+        model_config = model_type.config.from_dict(config)
+    model_config.kwargs = {"active_vocab_size": avs} if avs is not None else {}
+    args = CompileArgs(
+        model_config,
+        quantization,
+        model_type,
+        target,
+        opt,
+        build_func,
+        system_lib_prefix,
+        output,
+        overrides,
+        debug_dump,
+    )
+    args.display()
+    _compile(args, model_config)

python/mlc_llm/interface/compiler_flags.py

@@ -0,0 +1,227 @@
+"""Flags for overriding model config."""
+
+import dataclasses
+import enum
+from io import StringIO
+from typing import Optional
+
+from mlc_llm.support import argparse, logging
+from mlc_llm.support.config import ConfigOverrideBase
+
+logger = logging.getLogger(__name__)
+
+
+class IPCAllReduceStrategyType(enum.IntEnum):
+    """The all-reduce strategy."""
+
+    NONE = 0
+    ONESHOT = 1
+    TWOSHOT = 2
+    AUTO = 3
+
+
+@dataclasses.dataclass
+class OptimizationFlags:
+    """Optimization flags"""
+
+    flashinfer: bool = False
+    cublas_gemm: bool = False
+    faster_transformer: bool = False
+    cudagraph: bool = False
+    cutlass: bool = False
+    ipc_allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE
+
+    def __repr__(self) -> str:
+        out = StringIO()
+        print(f"flashinfer={int(self.flashinfer)}", file=out, end="")
+        print(f";cublas_gemm={int(self.cublas_gemm)}", file=out, end="")
+        print(f";faster_transformer={int(self.faster_transformer)}", file=out, end="")
+        print(f";cudagraph={int(self.cudagraph)}", file=out, end="")
+        print(f";cutlass={int(self.cutlass)}", file=out, end="")
+        print(
+            f";ipc_allreduce_strategy={self.ipc_allreduce_strategy.name}",
+            file=out,
+            end="",
+        )
+        return out.getvalue().rstrip()
+
+    @staticmethod
+    def from_str(source: str) -> "OptimizationFlags":
+        """Parse optimization flags from a string."""
+
+        if source in OPT_FLAG_PRESET:
+            return OPT_FLAG_PRESET[source]
+
+        def boolean(value: str) -> bool:
+            if value == "0":
+                return False
+            if value == "1":
+                return True
+            raise ValueError(f"Invalid boolean value: {value}")
+
+        parser = argparse.ArgumentParser(description="optimization flags")
+        parser.add_argument("--flashinfer", type=boolean, default=True)
+        parser.add_argument("--cublas_gemm", type=boolean, default=False)
+        parser.add_argument("--faster_transformer", type=boolean, default=False)
+        parser.add_argument("--cudagraph", type=boolean, default=False)
+        parser.add_argument("--cutlass", type=boolean, default=False)
+        parser.add_argument(
+            "--ipc_allreduce_strategy",
+            type=str,
+            choices=["NONE", "ONESHOT", "TWOSHOT", "AUTO"],
+            default="NONE",
+        )
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return OptimizationFlags(
+            flashinfer=results.flashinfer,
+            cublas_gemm=results.cublas_gemm,
+            faster_transformer=results.faster_transformer,
+            cudagraph=results.cudagraph,
+            cutlass=results.cutlass,
+            ipc_allreduce_strategy=IPCAllReduceStrategyType[results.ipc_allreduce_strategy],
+        )
+
+    def update(self, target, quantization) -> None:
+        """Update optimization flags based on additional information."""
+
+        def _flashinfer(target) -> bool:
+            from mlc_llm.support.auto_target import (
+                detect_cuda_arch_list,
+            )
+
+            if not self.flashinfer:
+                return False
+            if target.kind.name != "cuda":
+                return False
+            arch_list = detect_cuda_arch_list(target)
+            for arch in arch_list:
+                if arch < 80:
+                    logger.warning("flashinfer is not supported on CUDA arch < 80")
+                    return False
+            return True
+
+        def _cublas_gemm(target, quantization) -> bool:
+            """correct cublas_gemm flag"""
+            if target.kind.name not in ["cuda", "rocm"]:
+                return False
+            if not (
+                quantization.name in ["q0f16", "q0bf16", "q0f32"]
+                or "e4m3" in quantization.name
+                or "e5m2" in quantization.name
+            ):
+                return False
+            return self.cublas_gemm
+
+        def _faster_transformer(target) -> bool:
+            """correct faster_transformer flag"""
+            if not target.kind.name == "cuda":
+                return False
+            return self.faster_transformer
+
+        def _cutlass(target) -> bool:
+            """correct cutlass flag"""
+            if not target.kind.name == "cuda":
+                return False
+            return self.cutlass
+
+        def _cudagraph(target) -> bool:
+            """correct cudagraph flag"""
+            if not target.kind.name == "cuda":
+                return False
+            return self.cudagraph
+
+        self.flashinfer = _flashinfer(target)
+        self.cublas_gemm = _cublas_gemm(target, quantization)
+        self.faster_transformer = _faster_transformer(target)
+        self.cutlass = _cutlass(target)
+        self.cudagraph = _cudagraph(target)
+
+
+@dataclasses.dataclass
+class ModelConfigOverride(ConfigOverrideBase):
+    """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
+    max_batch_size: Optional[int] = None
+    tensor_parallel_shards: Optional[int] = None
+    pipeline_parallel_stages: Optional[int] = None
+    disaggregation: Optional[bool] = None
+
+    def __repr__(self) -> str:
+        out = StringIO()
+        print(f"context_window_size={self.context_window_size}", file=out, end="")
+        print(f";sliding_window_size={self.sliding_window_size}", file=out, end="")
+        print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="")
+        print(f";attention_sink_size={self.attention_sink_size}", file=out, end="")
+        print(f";max_batch_size={self.max_batch_size}", file=out, end="")
+        print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="")
+        print(
+            f";pipeline_parallel_stages={self.pipeline_parallel_stages}",
+            file=out,
+            end="",
+        )
+        print(f";disaggregation={self.disaggregation}", file=out, end="")
+        return out.getvalue().rstrip()
+
+    @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("--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)
+        parser.add_argument("--max_batch_size", type=int, default=None)
+        parser.add_argument("--tensor_parallel_shards", type=int, default=None)
+        parser.add_argument("--pipeline_parallel_stages", type=int, default=None)
+        parser.add_argument(
+            "--disaggregation",
+            type=lambda x: str(x).lower() in ["true", "1", "yes", "True"],
+            default=None,
+        )
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return ModelConfigOverride(
+            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,
+            max_batch_size=results.max_batch_size,
+            tensor_parallel_shards=results.tensor_parallel_shards,
+            pipeline_parallel_stages=results.pipeline_parallel_stages,
+            disaggregation=results.disaggregation,
+        )
+
+
+OPT_FLAG_PRESET = {
+    "O0": OptimizationFlags(
+        flashinfer=False,
+        cublas_gemm=False,
+        cudagraph=False,
+    ),
+    "O1": OptimizationFlags(
+        flashinfer=False,
+        cublas_gemm=True,
+        faster_transformer=True,
+        cudagraph=False,
+        cutlass=True,
+    ),
+    "O2": OptimizationFlags(
+        flashinfer=True,
+        cublas_gemm=True,
+        faster_transformer=False,
+        cudagraph=True,
+        cutlass=True,
+        ipc_allreduce_strategy=IPCAllReduceStrategyType.NONE,
+    ),
+    "O3": OptimizationFlags(
+        flashinfer=True,
+        cublas_gemm=True,
+        faster_transformer=True,
+        cudagraph=True,
+        cutlass=True,
+        ipc_allreduce_strategy=IPCAllReduceStrategyType.AUTO,
+    ),
+}

python/mlc_llm/interface/convert_weight.py

@@ -0,0 +1,250 @@
+"""Python entrypoint of weight conversion."""
+
+import contextlib
+import dataclasses
+import math
+import os
+import tempfile
+from collections.abc import Iterator
+from io import StringIO
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple  # noqa: UP035
+
+from tvm import tirx
+from tvm.contrib import tvmjs
+from tvm.runtime import DataType, Device, Tensor
+from tvm.runtime import cpu as cpu_device
+from tvm.target import Target
+
+from mlc_llm.loader import LOADER
+from mlc_llm.model import Model
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import logging, tqdm
+from mlc_llm.support.auto_weight import detect_weight
+from mlc_llm.support.preshard import apply_preshard
+from mlc_llm.support.style import bold, green
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class ConversionArgs:
+    """Arguments to MLC LLM's weight conversation and quantization flow."""
+
+    config: Path
+    quantization: Quantization
+    model: Model
+    device: Device
+    source: Path
+    source_format: str
+    output: Path
+    lora_adapter: Optional[Path] = None
+
+    def display(self) -> None:
+        """Display the arguments to stdout."""
+
+        def _device_to_str(device: Device) -> str:
+            return f"{Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}"
+
+        out = StringIO()
+        print(f"{bold('Weight conversion with arguments:')}", file=out)
+        print(f"  {bold('--config'):<25} {self.config}", file=out)
+        print(f"  {bold('--quantization'):<25} {self.quantization}", file=out)
+        print(f"  {bold('--model-type'):<25} {self.model.name}", file=out)
+        print(f"  {bold('--device'):<25} {_device_to_str(self.device)}", file=out)
+        print(f"  {bold('--source'):<25} {self.source}", file=out)
+        print(f"  {bold('--source-format'):<25} {self.source_format}", file=out)
+        print(f"  {bold('--output'):<25} {self.output}", file=out)
+        if self.lora_adapter is not None:
+            print(f"  {bold('--lora-adapter'):<25} {self.lora_adapter}", file=out)
+        print(out.getvalue().rstrip())
+
+
+def _resolve_base_model_dir(source: Path) -> Path:
+    return source if source.is_dir() else source.parent
+
+
+@contextlib.contextmanager
+def _merge_lora_adapter_with_base_model(base_source: Path, lora_adapter: Path) -> Iterator[Path]:
+    base_model_dir = _resolve_base_model_dir(base_source)
+    if not base_model_dir.exists():
+        raise ValueError(f"Base model directory does not exist: {base_model_dir}")
+    if not lora_adapter.exists() or not lora_adapter.is_dir():
+        raise ValueError(f"LoRA adapter directory does not exist: {lora_adapter}")
+
+    try:
+        from peft import PeftModel
+        from transformers import AutoModelForCausalLM
+
+    except ImportError as err:
+        raise ImportError(
+            "`--lora-adapter` requires `peft` and `transformers` to be installed."
+        ) from err
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        merged_model_dir = Path(temp_dir) / "merged_model"
+        logger.info("Merging LoRA adapter %s into base model %s", lora_adapter, base_model_dir)
+
+        base_model = AutoModelForCausalLM.from_pretrained(
+            str(base_model_dir),
+            torch_dtype="auto",
+            trust_remote_code=False,
+            low_cpu_mem_usage=True,
+        )
+        merged_model = PeftModel.from_pretrained(
+            base_model, str(lora_adapter), is_trainable=False
+        ).merge_and_unload()
+        merged_model.save_pretrained(str(merged_model_dir), safe_serialization=True)
+        yield merged_model_dir
+
+
+def _convert_args(args: ConversionArgs) -> None:
+    pre_shards_num = os.getenv("MLC_INTERNAL_PRESHARD_NUM")
+    # model config & quantization config
+    model_config = args.model.config.from_file(args.config)
+    if (
+        args.quantization.kind == "ft-quant"
+        and hasattr(model_config, "tensor_parallel_shards")
+        and model_config.tensor_parallel_shards > 1
+    ):
+        raise NotImplementedError
+    if pre_shards_num is not None:
+        model_config.tensor_parallel_shards = int(pre_shards_num)
+    model, quantize_map = args.model.quantize[args.quantization.kind](
+        model_config, args.quantization
+    )
+    _, _named_params, _ = model.export_tvm(
+        spec=model.get_default_spec(),
+        allow_extern=True,
+    )
+    named_params = dict(_named_params)
+
+    if pre_shards_num is not None:
+        named_params, preshard_funcs = apply_preshard(named_params, int(pre_shards_num), args)
+    else:
+        preshard_funcs = None
+
+    def _check_param(name: str, param: Tensor):
+        nonlocal named_params
+        if name not in named_params:
+            raise ValueError(f"Parameter not found in model: {name}")
+        if name in param_names:
+            raise ValueError(f"Duplication: Parameter {name} already computed")
+
+        # Check shape (possibly dynamic)
+        def _check_shape(actual: tuple, expect: tuple):  # expect can have tirx.Var
+            if len(actual) != len(expect):
+                return False
+            for actual_i, expect_i in zip(actual, expect):
+                assert isinstance(expect_i, (int, tirx.Var))
+                if isinstance(expect_i, int) and actual_i != expect_i:
+                    return False
+            return True
+
+        expect_shape = named_params[name].shape
+        actual_shape = param.shape
+        if not _check_shape(actual_shape, expect_shape):
+            raise ValueError(
+                f"Parameter {name} has shape {param.shape}, but expected {expect_shape}"
+            )
+        # Check dtype
+        actual_dtype = param.dtype
+        expect_dtype = named_params[name].dtype
+        if actual_dtype != expect_dtype:
+            raise ValueError(
+                f"Parameter {name} has dtype {param.dtype}, but expected {expect_dtype}"
+            )
+        del named_params[name]
+
+    # load and quantize
+    param_names = set()
+    total_bytes = 0.0
+    total_params: int = 0
+
+    def _param_generator() -> Iterator[Tuple[str, Tensor]]:  # noqa: UP006
+        nonlocal total_params, total_bytes
+        with Target.from_device(args.device), tqdm.redirect():
+            loader = LOADER[args.source_format](
+                path=args.source,
+                extern_param_map=args.model.source[args.source_format](
+                    model_config, args.quantization
+                ),
+                quantize_param_map=quantize_map,
+            )
+            for name, param in loader.load(device=args.device, preshard_funcs=preshard_funcs):
+                _check_param(name, param)
+                param_names.add(name)
+                param = param.copyto(cpu_device())
+                total_bytes += math.prod(param.shape) * DataType(param.dtype).itemsize
+                yield name, param
+        total_params = loader.stats.total_param_num
+
+    def _metadata_callback() -> Dict[str, Any]:  # noqa: UP006
+        return {
+            "ParamSize": len(param_names),
+            "ParamBytes": total_bytes,
+            "BitsPerParam": total_bytes * 8.0 / total_params,
+        }
+
+    # dump to output directory
+    tvmjs.dump_tensor_cache(
+        _param_generator(),
+        str(args.output),
+        meta_data=_metadata_callback,
+        encode_format="f32-to-bf16",
+        show_progress=False,
+    )
+    if named_params:
+        raise ValueError(f"Parameter not found in source: {', '.join(named_params.keys())}")
+    # Log necessary statistics
+    logger.info(
+        "%s after quantization: %.3f GB",
+        green("Parameter size"),
+        total_bytes / (1024**3),
+    )
+    logger.info(f"%s: {total_params:,}", green("Total parameters"))
+    logger.info(
+        "%s: %.3f",
+        green("Bits per parameter"),
+        total_bytes * 8.0 / total_params,
+    )
+    logger.info("Saved to directory: %s", bold(str(args.output)))
+
+
+def convert_weight(
+    config: Path,
+    quantization: Quantization,
+    model: Model,
+    device: Device,
+    source: Path,
+    source_format: str,
+    output: Path,
+    lora_adapter: Optional[Path] = None,
+):
+    """MLC LLM's weight conversation and quantization flow."""
+    args = ConversionArgs(
+        config, quantization, model, device, source, source_format, output, lora_adapter
+    )
+
+    allowed_lora_source_formats = {"huggingface-safetensor", "huggingface-torch"}
+    if lora_adapter is not None and source_format not in allowed_lora_source_formats:
+        raise ValueError(
+            f"`--lora-adapter` only supports source formats: {sorted(allowed_lora_source_formats)}"
+        )
+
+    if lora_adapter is not None:
+        with _merge_lora_adapter_with_base_model(source, lora_adapter) as merged_model_dir:
+            merged_source, merged_source_format = detect_weight(
+                weight_path=merged_model_dir,
+                config_json_path=config,
+                weight_format="auto",
+            )
+            merged_args = dataclasses.replace(
+                args, source=merged_source, source_format=merged_source_format
+            )
+            merged_args.display()
+            _convert_args(merged_args)
+            return
+
+    args.display()
+    _convert_args(args)

python/mlc_llm/interface/gen_config.py

@@ -0,0 +1,359 @@
+"""Generator of mlc-chat-config.json and tokenizer configuration."""
+
+import dataclasses
+import json
+import re
+import shutil
+from dataclasses import asdict
+from pathlib import Path
+from typing import Optional
+
+from mlc_llm.conversation_template import ConvTemplateRegistry
+from mlc_llm.model import Model
+from mlc_llm.protocol.mlc_chat_config import MLCChatConfig
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import convert_tiktoken, logging
+from mlc_llm.support.style import bold, green, red
+from mlc_llm.tokenizers import Tokenizer
+
+from .compiler_flags import ModelConfigOverride
+
+logger = logging.getLogger(__name__)
+
+FOUND = green("Found")
+NOT_FOUND = red("Not found")
+FAILED = red("Failed")
+
+
+def apply_system_defaults_for_missing_fields(mlc_chat_config: MLCChatConfig) -> None:
+    """Apply system default value."""
+    for key, value in mlc_chat_config.get_system_defaults_for_missing_fields().items():
+        setattr(mlc_chat_config, key, value)
+        logger.info("[System default] Setting %s: %s", bold(key), value)
+
+
+def check_string(s: str) -> bool:
+    """Check whether it's a string."""
+    s = s[1:] if s[0] == "b" else s
+    delimit = s[0]
+    if s[-1] != delimit or delimit not in ["'", '"']:
+        return False
+    for i in range(1, len(s) - 1):
+        if s[i] == delimit and s[i - 1] != "\\":
+            return False
+    return True
+
+
+def txt2rwkv_tokenizer(vocab: Path, out: Path) -> None:
+    """Generate tokenizer_model from RWKV vocab file."""
+    idx2token = {}
+
+    with vocab.open("r", encoding="utf-8") as f:
+        lines = f.readlines()
+
+    for line in lines:
+        idx = int(line[: line.index(" ")])
+        raw = line[line.index(" ") : line.rindex(" ")].strip()
+        if check_string(raw):
+            x = eval(raw)
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+            assert len(x) == int(line[line.rindex(" ") :])
+            idx2token[idx] = x
+        else:
+            raise ValueError("Unsupported vocab dictionary")
+
+    with (out / "tokenizer_model").open("wb") as f:
+        import msgpack
+
+        msgpack.pack(idx2token, f)
+
+
+def json2rwkv_tokenizer(vocab: Path, out: Path) -> None:
+    """Generate tokenizer_model from RWKV vocab file."""
+    idx2token = {}
+
+    with vocab.open("r", encoding="utf-8") as f:
+        data = json.load(f)
+        for key, value in data.items():
+            x = key.encode("utf-8") if isinstance(key, str) else key
+            assert isinstance(x, bytes)
+            idx2token[int(value)] = x
+
+    with (out / "tokenizer_model").open("wb") as f:
+        import msgpack
+
+        msgpack.pack(idx2token, f)
+
+
+def gen_config(
+    config: Path,
+    model: Model,
+    quantization: Quantization,
+    conv_template: str,
+    context_window_size: Optional[int],
+    sliding_window_size: Optional[int],
+    prefill_chunk_size: Optional[int],
+    attention_sink_size: Optional[int],
+    tensor_parallel_shards: Optional[int],
+    pipeline_parallel_stages: Optional[int],
+    disaggregation: Optional[bool],
+    max_batch_size: int,
+    output: Path,
+):
+    """Entrypoint of MLC Chat configuration generation."""
+    # Step 1. Initialize `mlc-chat-config.json` using `config.json`
+    conversation_reg = ConvTemplateRegistry.get_conv_template(conv_template)
+    if conversation_reg is None:
+        logger.warning(
+            "%s: Conversation template is not registered in ConvTemplateRegistry: %s",
+            red("Warning"),
+            conv_template,
+        )
+        conversation = conv_template
+    else:
+        conversation = conversation_reg.to_json_dict()
+
+    model_config = ModelConfigOverride(
+        context_window_size=context_window_size,
+        sliding_window_size=sliding_window_size,
+        prefill_chunk_size=prefill_chunk_size,
+        attention_sink_size=attention_sink_size,
+        max_batch_size=max_batch_size,
+        tensor_parallel_shards=tensor_parallel_shards,
+        pipeline_parallel_stages=pipeline_parallel_stages,
+        disaggregation=disaggregation,
+    ).apply(model.config.from_file(config))
+    mlc_chat_config = MLCChatConfig(
+        model_type=model.name,
+        quantization=quantization.name,
+        model_config=model_config.asdict(),
+        vocab_size=model_config.vocab_size,
+        active_vocab_size=getattr(model_config, "active_vocab_size", model_config.vocab_size),
+        context_window_size=getattr(model_config, "context_window_size", -1),
+        sliding_window_size=getattr(model_config, "sliding_window_size", -1),
+        prefill_chunk_size=model_config.prefill_chunk_size,
+        attention_sink_size=getattr(model_config, "attention_sink_size", -1),
+        tensor_parallel_shards=model_config.tensor_parallel_shards,
+        pipeline_parallel_stages=getattr(model_config, "pipeline_parallel_stages", 1),
+        disaggregation=getattr(model_config, "disaggregation", False),
+        conv_template=conversation,
+        model_task=model.model_task,
+        embedding_metadata=(
+            dataclasses.asdict(model.embedding_metadata) if model.embedding_metadata else None
+        ),
+    )
+    # Step 2. Load `generation_config.json` and `config.json` for text-generation related configs
+    for generation_config_filename in ["generation_config.json", "config.json"]:
+        generation_config = config.parent / generation_config_filename
+        if generation_config.exists():
+            with generation_config.open("r", encoding="utf-8") as in_file:
+                generation_config_json = json.load(in_file)
+            for key, value in generation_config_json.items():
+                if hasattr(mlc_chat_config, key) and getattr(mlc_chat_config, key) is None:
+                    setattr(mlc_chat_config, key, value)
+                    logger.info(
+                        "[%s] Setting %s: %s",
+                        generation_config_filename,
+                        bold(key),
+                        value,
+                    )
+        else:
+            logger.info("%s %s: %s", NOT_FOUND, generation_config_filename, generation_config)
+
+    # Step 3. Copy tokenizer configuration
+    # 3.1. Copy over the files and populate mlc_chat_config
+    for filename in TOKENIZER_FILES:
+        file = config.parent / filename
+        if file.exists():
+            mlc_chat_config.tokenizer_files.append(filename)
+            dest = output / filename
+            shutil.copy(file, dest)
+            logger.info("%s tokenizer config: %s. Copying to %s", FOUND, file, bold(str(dest)))
+        else:
+            logger.info("%s tokenizer config: %s", NOT_FOUND, file)
+    # 3.2. Generate `tokenizer_model` for rwkv if `rwkv_vocab_.*` is found
+    pattern = re.compile(r"rwkv_vocab_v\d{8}\.(json|txt)")
+    for item in config.parent.iterdir():
+        if item.is_file() and pattern.match(item.name):
+            logger.info(
+                "%s RWKV vocab file: %s. Genetating %s",
+                FOUND,
+                item,
+                bold("tokenizer_model"),
+            )
+            if item.name.endswith(".txt"):
+                txt2rwkv_tokenizer(item, output)
+            else:
+                json2rwkv_tokenizer(item, output)
+    # 3.3. If we have `tokenizer.model` but not `tokenizer.json`, try convert it to
+    # `tokenizer.json` with `transformers`.
+    tokenizer_json_file = config.parent / "tokenizer.json"
+    tokenizer_model_file = config.parent / "tokenizer.model"
+    if tokenizer_model_file.exists() and (not tokenizer_json_file.exists()):
+        logger.info(
+            "The model has `tokenizer.model` but not `tokenizer.json`. "
+            "It is always recommended to prefer JSON instead. "
+            "Attempting to convert using HuggingFace transformers library"
+        )
+        try:
+            from transformers import (
+                AutoTokenizer,
+            )
+
+            tokenizer_json_save_dest = output / "tokenizer.json"
+            fast_tokenizer = AutoTokenizer.from_pretrained(str(config.parent), use_fast=True)
+            fast_tokenizer.backend_tokenizer.save(str(tokenizer_json_save_dest))
+            mlc_chat_config.tokenizer_files.append("tokenizer.json")
+            logger.info(
+                "Successfully converted `tokenizer.model` to: %s",
+                tokenizer_json_save_dest,
+            )
+        except Exception:
+            logger.warning(
+                "Converting to `tokenizer.json` %s with the exception below. "
+                "Skipping the conversion.",
+                FAILED,
+                exc_info=True,
+            )
+    # 3.3. If we still don't have "tokenizer.json" at this point, try looking for "*.tiktoken" files
+    if (not tokenizer_json_file.exists()) and list(config.parent.glob("*.tiktoken")):
+        try:
+            logger.info(
+                "The model has tiktoken files but not `tokenizer.json`. "
+                "Attempting to convert from tiktoken files"
+            )
+            convert_tiktoken.convert_tiktoken(
+                str(config.parent), str(output), mlc_chat_config.context_window_size
+            )
+            mlc_chat_config.tokenizer_files.append("tokenizer.json")
+            mlc_chat_config.tokenizer_files.append("vocab.json")
+            mlc_chat_config.tokenizer_files.append("merges.txt")
+            mlc_chat_config.tokenizer_files.append("special_tokens_map.json")
+            logger.info("Succesfully converted from tiktoken files to: %s", str(output))
+        except Exception:
+            logger.exception("%s with the exception below. Skipping", FAILED)
+
+    # 3.4. Detect tokenizer info
+    mlc_chat_config.tokenizer_info = asdict(Tokenizer.detect_tokenizer_info(str(output)))
+    logger.info("Detected tokenizer info: %s", mlc_chat_config.tokenizer_info)
+
+    # 3.5. Ensure added_tokens do not have duplicated added_tokens, a mistake from model releaser
+    # that affects correctness of huggingface tokenizer.
+    # See https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B/discussions/15.
+    if tokenizer_json_file.exists():
+        with open(tokenizer_json_file, encoding="utf-8") as f:
+            tokenizer_json = json.load(f)
+            if "added_tokens" in tokenizer_json:
+                appeared_content = set()
+                for added_token in tokenizer_json["added_tokens"]:
+                    content = added_token["content"]
+                    if content in appeared_content:
+                        logger.exception(
+                            "%s with incorrect tokenizer.json which has duplicated token %s. "
+                            "This affects correctness of huggingface tokenizer during runtime, "
+                            "please check your tokenizer.json to remove duplication manually.",
+                            FAILED,
+                            content,
+                        )
+                        raise ValueError("Duplicated vocab in tokenizer.json")
+                    appeared_content.add(content)
+
+    # Step 4. Load system default value
+    apply_system_defaults_for_missing_fields(mlc_chat_config)
+
+    # Step 5. Use HF tokenizer to detect active vocab size via len(tokenizer)
+    if tokenizer_json_file.exists():
+        try:
+            from transformers import (
+                AutoTokenizer,
+            )
+
+            hf_tokenizer = AutoTokenizer.from_pretrained(str(config.parent), use_fast=True)
+            active_vocab_size = len(hf_tokenizer)
+            if mlc_chat_config.active_vocab_size != active_vocab_size:
+                logger.info(
+                    "Overriding active_vocab_size from %d to %d using HF tokenizer",
+                    mlc_chat_config.active_vocab_size,
+                    active_vocab_size,
+                )
+                mlc_chat_config.active_vocab_size = active_vocab_size
+        except Exception:
+            logger.warning(
+                "Detecting active_vocab_size %s with the exception below. Skipping.",
+                FAILED,
+                exc_info=True,
+            )
+
+    # Step 5. Dump the configuration file to output directory
+    with (output / "mlc-chat-config.json").open("w", encoding="utf-8") as out_file:
+        json.dump(mlc_chat_config.model_dump(by_alias=True), out_file, indent=2)
+        logger.info("Dumping configuration file to: %s", bold(out_file.name))
+
+
+TOKENIZER_FILES = [
+    "tokenizer.model",
+    "tokenizer.json",
+    "vocab.json",
+    "merges.txt",
+    "added_tokens.json",
+    "tokenizer_config.json",
+]
+# FIXME: Copy RWKV tokenizer file
+
+CONV_TEMPLATES = {
+    "llama-4",
+    "llama-3",
+    "llama-3_1",
+    "chatml",
+    "chatml_nosystem",
+    "qwen2",
+    "open_hermes_mistral",
+    "neural_hermes_mistral",
+    "llama_default",
+    "llama-2",
+    "mistral_default",
+    "ministral3",
+    "ministral3_reasoning",
+    "gpt2",
+    "codellama_completion",
+    "codellama_instruct",
+    "redpajama_chat",
+    "rwkv_world",
+    "gorilla",
+    "gorilla-openfunctions-v2",
+    "dolly",
+    "oasst",
+    "stablelm",
+    "LM",
+    "stablelm-3b",
+    "gpt_bigcode",
+    "wizardlm_7b",
+    "wizard_coder_or_math",
+    "glm",
+    "phi-2",
+    "phi-3",
+    "phi-3-vision",
+    "phi-4",
+    "stablelm-2",
+    "gemma_instruction",
+    "gemma3_instruction",
+    "orion",
+    "llava",
+    "hermes2_pro_llama3",
+    "hermes3_llama-3_1",
+    "tinyllama_v1_0",
+    "aya-23",
+    "deepseek",
+    "deepseek_v2",
+    "deepseek_v3",
+    "deepseek_r1_qwen",
+    "deepseek_r1_llama",
+    "olmo",
+    "olmo2",
+    "nemotron",
+    "llm-jp",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_5_nothink",
+}

python/mlc_llm/interface/help.py

@@ -0,0 +1,273 @@
+"""Help message for CLI arguments."""
+
+HELP = {
+    "config": (
+        """
+1) Path to a HuggingFace model directory that contains a `config.json` or
+2) Path to `config.json` in HuggingFace format, or
+3) The name of a pre-defined model architecture.
+
+A `config.json` file in HuggingFace format defines the model architecture, including the vocabulary
+size, the number of layers, the hidden size, number of attention heads, etc.
+Example: https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json.
+
+A HuggingFace directory often contains a `config.json` which defines the model architecture,
+the non-quantized model weights in PyTorch or SafeTensor format, tokenizer configurations,
+as well as an optional `generation_config.json` provides additional default configuration for
+text generation.
+Example: https://huggingface.co/codellama/CodeLlama-7b-hf/tree/main.
+"""
+    ).strip(),
+    "quantization": """
+The quantization mode we use to compile. If unprovided, will infer from `model`.
+""".strip(),
+    "model": """
+A path to ``mlc-chat-config.json``, or an MLC model directory that contains `mlc-chat-config.json`.
+It can also be a link to a HF repository pointing to an MLC compiled model.
+""".strip(),
+    "model_lib": """
+The full path to the model library file to use (e.g. a ``.so`` file). If unspecified, we will use
+the provided ``model`` to search over possible paths. It the model lib is not found, it will be
+compiled in a JIT manner.
+""".strip(),
+    "model_type": """
+Model architecture such as "llama". If not set, it is inferred from `mlc-chat-config.json`.
+""".strip(),
+    "device_compile": """
+The GPU device to compile the model to. If not set, it is inferred from GPUs available locally.
+""".strip(),
+    "enable_subgroups": """
+Enable WebGPU subgroups in codegen. This only applies to WebGPU targets and will set
+supports_subgroups accordingly.
+""".strip(),
+    "device_quantize": """
+The device used to do quantization such as "cuda" or "cuda:0". Will detect from local available GPUs
+if not specified.
+""".strip(),
+    "device_deploy": """
+The device used to deploy the model such as "cuda" or "cuda:0". Will detect from local
+available GPUs if not specified.
+""".strip(),
+    "host": """
+The host LLVM triple to compile the model to. If not set, it is inferred from the local CPU and OS.
+Examples of the LLVM triple:
+1) iPhones: arm64-apple-ios;
+2) ARM64 Android phones: aarch64-linux-android;
+3) WebAssembly: wasm32-unknown-unknown-wasm;
+4) Windows: x86_64-pc-windows-msvc;
+5) ARM macOS: arm64-apple-darwin.
+""".strip(),
+    "opt": """
+Optimization flags. MLC LLM maintains a predefined set of optimization flags,
+denoted as O0, O1, O2, O3, where O0 means no optimization, O2 means majority of them,
+and O3 represents extreme optimization that could potentially break the system.
+Meanwhile, optimization flags could be explicitly specified via details knobs, e.g.
+--opt="cublas_gemm=1;cudagraph=0".
+""".strip(),
+    "system_lib_prefix": """
+Adding a prefix to all symbols exported. Similar to "objcopy --prefix-symbols".
+This is useful when compiling multiple models into a single library to avoid symbol
+conflicts. Different from objcopy, this takes no effect for shared library.
+""".strip(),
+    "context_window_size": """
+Option to provide the maximum sequence length supported by the model.
+This is usually explicitly shown as context length or context window in the model card.
+If this option is not set explicitly, by default,
+it will be determined by `context_window_size` or `max_position_embeddings` in `config.json`,
+and the latter is usually inaccurate for some models.
+""".strip(),
+    "output_compile": """
+The path to the output file. The suffix determines if the output file is a shared library or
+objects. Available suffixes:
+1) Linux: .so (shared), .tar (objects);
+2) macOS: .dylib (shared), .tar (objects);
+3) Windows: .dll (shared), .tar (objects);
+4) Android, iOS: .tar (objects);
+5) Web: .wasm (web assembly).
+""".strip(),
+    "source": """
+The path to original model weight, infer from `config` if missing.
+""".strip(),
+    "source_format": """
+The format of source model weight, infer from `config` if missing.
+""".strip(),
+    "output_quantize": """
+The output directory to save the quantized model weight. Will create `params_shard_*.bin` and
+`tensor-cache.json` in this directory.
+""".strip(),
+    "conv_template": """
+Conversation template. It depends on how the model is tuned. Use "LM" for vanilla base model
+""".strip(),
+    "output_gen_mlc_chat_config": """
+The output directory for generated configurations, including `mlc-chat-config.json` and tokenizer
+configuration.
+""".strip(),
+    "sliding_window_size": """
+(Experimental) The sliding window size in sliding window attention (SWA).
+This optional field overrides the `sliding_window_size` in config.json for
+those models that use SWA. Currently only useful when compiling Mistral.
+This flag subjects to future refactoring.
+""".strip(),
+    "prefill_chunk_size": """
+(Experimental) The chunk size during prefilling. By default,
+the chunk size is the same as sliding window or max sequence length.
+This flag subjects to future refactoring.
+""".strip(),
+    "attention_sink_size": """
+(Experimental) The number of stored sinks. Only supported on Mistral yet. By default,
+the number of sinks is 4. This flag subjects to future refactoring.
+""".strip(),
+    "max_batch_size": """
+The maximum allowed batch size set for the KV cache to concurrently support.
+""".strip(),
+    """tensor_parallel_shards""": """
+Number of shards to split the model into in tensor parallelism multi-gpu inference.
+""".strip(),
+    """pipeline_parallel_stages""": """
+Number of pipeline stages to split the model layers for pipeline parallelism.
+""".strip(),
+    """disaggregation""": """
+Whether enable disaggregation when compiling the model.
+""".strip(),
+    "overrides": """
+Model configuration override. Configurations to override `mlc-chat-config.json`. Supports
+`context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`,
+`max_batch_size` and `tensor_parallel_shards`. Meanwhile, model config could be explicitly
+specified via details knobs, e.g. --overrides "context_window_size=1024;prefill_chunk_size=128".
+""".strip(),
+    "modelconfig_overrides": """
+Model configuration override. Supports overriding,
+`context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`,
+`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": """
+Specifies the directory where the compiler will store its IRs for debugging purposes
+during various phases of compilation. By default, this is set to `None`, indicating
+that debug dumping is disabled.
+""".strip(),
+    "prompt": """
+The prompt of the text generation.
+""".strip(),
+    "generate_length": """
+The target length of the text generation.
+""".strip(),
+    "max_total_sequence_length_serve": """
+The KV cache total token capacity, i.e., the maximum total number of tokens that
+the KV cache support. This decides the GPU memory size that the KV cache consumes.
+If not specified, system will automatically estimate the maximum capacity based
+on the vRAM size on GPU.
+""".strip(),
+    "prefill_chunk_size_serve": """
+The maximum number of tokens the model passes for prefill each time.
+It should not exceed the prefill chunk size in model config.
+If not specified, this defaults to the prefill chunk size in model config.
+""".strip(),
+    "max_history_size_serve": """
+The maximum history length for rolling back the RNN state.
+If unspecified, the default value is 1.
+KV cache does not need this.
+""".strip(),
+    "enable_tracing_serve": """
+Enable Chrome Tracing for the server.
+After enabling, you can send POST request to the "debug/dump_event_trace" entrypoint
+to get the Chrome Trace. For example,
+"curl -X POST http://127.0.0.1:8000/debug/dump_event_trace -H "Content-Type: application/json" -d '{"model": "dist/llama"}'"
+""".strip(),  # noqa: E501
+    "mode_serve": """
+The engine mode in MLC LLM. We provide three preset modes: "local", "interactive" and "server".
+The default mode is "local".
+The choice of mode decides the values of "max_num_sequence", "max_total_seq_length" and
+"prefill_chunk_size" when they are not explicitly specified.
+1. Mode "local" refers to the local server deployment which has low request concurrency.
+   So the max batch size will be set to 4, and max total sequence length and prefill chunk size
+   are set to the context window size (or sliding window size) of the model.
+2. Mode "interactive" refers to the interactive use of server, which has at most 1 concurrent
+   request. So the max batch size will be set to 1, and max total sequence length and prefill
+   chunk size are set to the context window size (or sliding window size) of the model.
+3. Mode "server" refers to the large server use case which may handle many concurrent request
+   and want to use GPU memory as much as possible. In this mode, we will automatically infer
+   the largest possible max batch size and max total sequence length.
+You can manually specify arguments "max_num_sequence", "max_total_seq_length" and
+"prefill_chunk_size" via "--overrides" to override the automatic inferred values.
+For example: --overrides "max_num_sequence=32;max_total_seq_length=4096"
+""".strip(),
+    "additional_models_serve": """
+The model paths and (optional) model library paths of additional models (other than the main model).
+When engine is enabled with speculative decoding, additional models are needed.
+The way of specifying additional models is:
+"--additional-models model_path_1 model_path_2 ..." or
+"--additional-models model_path_1,model_lib_1 model_path_2 ...".
+When the model lib of a model is not given, JIT model compilation will be activated
+to compile the model automatically.
+""".strip(),
+    "gpu_memory_utilization_serve": """
+A number in (0, 1) denoting the fraction of GPU memory used by the server in total.
+It is used to infer to maximum possible KV cache capacity.
+When it is unspecified, it defaults to 0.85.
+Under mode "local" or "interactive", the actual memory usage may be significantly smaller than
+this number. Under mode "server", the actual memory usage may be slightly larger than this number.
+""".strip(),
+    "speculative_mode_serve": """
+The speculative decoding mode. Right now four options are supported:
+ - "disable", where speculative decoding is not enabled,
+ - "small_draft", denoting the normal speculative decoding (small draft) style,
+ - "eagle", denoting the eagle-style speculative decoding.
+ - "medusa", denoting the medusa-style speculative decoding.
+The default mode is "disable".
+""".strip(),
+    "spec_draft_length_serve": """
+The number of draft tokens to generate in speculative proposal.
+Being 0 means to enable adaptive speculative mode, where the draft length will be
+automatically adjusted based on engine state. The default values is 0.
+""".strip(),
+    "prefix_cache_mode_serve": """
+The prefix cache mode. Right now two options are supported:
+ - "disable", where prefix cache is not enabled,
+ - "radix", denoting the normal paged radix tree based prefix cache,
+The default mode is "radix".
+""".strip(),
+    "prefix_cache_max_num_recycling_seqs_serve": """
+The maximum number of sequences in prefix cache, default as max_batch_size.
+And set 0 to disable prefix cache, set -1 to have infinite capacity prefix cache.
+""".strip(),
+    "prefill_mode": """
+The prefill mode. "chunked" means the basic prefill with chunked input enabled. "hybrid" means the
+hybrid prefill or split-fuse, so that decode step will be converted into prefill.
+""".strip(),
+    "overrides_serve": """
+Overriding extra configurable fields of EngineConfig and model compilation config.
+Supporting fields that can be be overridden: "tensor_parallel_shards", "max_num_sequence",
+"max_total_seq_length", "prefill_chunk_size", "max_history_size", "gpu_memory_utilization",
+"spec_draft_length", "prefix_cache_max_num_recycling_seqs", "context_window_size",
+"sliding_window_size", "attention_sink_size".
+Please check out the documentation of EngineConfig in mlc_llm/serve/config.py for detailed docstring
+of each field.
+Example: --overrides "max_num_sequence=32;max_total_seq_length=4096;tensor_parallel_shards=2"
+""".strip(),
+    "config_package": """
+The path to "mlc-package-config.json" which is used for package build.
+See "https://github.com/mlc-ai/mlc-llm/blob/main/ios/MLCChat/mlc-package-config.json" as an example.
+""".strip(),
+    "mlc_llm_source_dir": """
+The source code path to MLC LLM.
+""".strip(),
+    "output_package": """
+The path of output directory for the package build outputs.
+""".strip(),
+    "calibration_dataset": """
+The path to the calibration dataset.
+    """.strip(),
+    "num_calibration_samples": """
+The number of samples used for calibration.
+    """.strip(),
+    "output_calibration": """
+The output directory to save the calibration params.
+    """.strip(),
+    "seed_calibrate": """
+The seed to sample the calibration dataset.""",
+    "pd_balance_factor": """
+How much prefill to move to decode engine. For example,
+0.1 means the last 10 percent tokens are prefilled by decode engine.
+    """.strip(),
+}

python/mlc_llm/interface/jit.py

@@ -0,0 +1,181 @@
+"""Just-in-time compilation of MLC-Chat models."""
+
+import dataclasses
+import hashlib
+import json
+import os
+import shlex
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Optional, Union  # noqa: UP035
+
+from tvm.runtime import Device
+
+from mlc_llm.model import MODELS
+from mlc_llm.support import logging
+from mlc_llm.support.auto_device import device2str
+from mlc_llm.support.constants import (
+    MLC_DSO_SUFFIX,
+    MLC_JIT_POLICY,
+    MLC_LLM_HOME,
+    MLC_TEMP_DIR,
+)
+from mlc_llm.support.style import blue, bold
+
+from .compiler_flags import ModelConfigOverride, OptimizationFlags
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class JITResult:
+    """The jit compilation result class."""
+
+    model_lib_path: str
+    system_lib_prefix: Optional[str] = None
+
+
+def log_jit_policy():
+    """log current jit policy"""
+    logger.info(
+        "%s = %s. Can be one of: ON, OFF, REDO, READONLY",
+        bold("MLC_JIT_POLICY"),
+        MLC_JIT_POLICY,
+    )
+
+
+def jit(
+    model_path: Path,
+    overrides: Dict[str, Any],  # noqa: UP006
+    device: Union[Device, str],
+    system_lib_prefix: Optional[str] = None,
+    *,
+    skip_log_jit_policy=False,
+) -> JITResult:
+    """Just-in-time compile a MLC-Chat model."""
+    # skip logging jit policy since when outside can hint once
+    if not skip_log_jit_policy:
+        log_jit_policy()
+
+    if MLC_JIT_POLICY == "OFF":
+        raise RuntimeError("JIT is disabled by MLC_JIT_POLICY=OFF")
+
+    with open(model_path / "mlc-chat-config.json", encoding="utf-8") as in_file:
+        mlc_chat_config = json.load(in_file)
+    model_type = mlc_chat_config.pop("model_type")
+    quantization = mlc_chat_config.pop("quantization")
+    lib_suffix = MLC_DSO_SUFFIX if device not in ["iphone", "macabi", "android"] else "tar"
+
+    def _get_optimization_flags() -> str:
+        opt = overrides.pop("opt", None)
+        if opt is None:
+            opt = "O2"
+        return repr(OptimizationFlags.from_str(opt))
+
+    def _get_overrides() -> str:
+        forbid_list = [
+            "context_window_size",
+            "sliding_window_size",
+            "attention_sink_size",
+        ]
+        result = []
+        for field in dataclasses.fields(ModelConfigOverride):
+            value = overrides.get(field.name, None)
+            if value is not None:
+                if field.name in forbid_list and value == -1:
+                    continue
+                result.append(f"{field.name}={value}")
+        return ";".join(result)
+
+    def _get_model_config() -> Dict[str, Any]:  # noqa: UP006
+        model_config = mlc_chat_config.pop("model_config")
+        model_config.update(mlc_chat_config)
+        for field in dataclasses.fields(ModelConfigOverride):
+            value = overrides.get(field.name, None)
+            if value is not None:
+                model_config[field.name] = value
+        return MODELS[model_type].config.from_dict(model_config).asdict()
+
+    def _run_jit(
+        opt: str,
+        overrides: str,
+        device: str,
+        system_lib_prefix: Optional[str],
+        dst: str,
+    ):
+        with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir:
+            dso_path = os.path.join(tmp_dir, f"lib.{lib_suffix}")
+            cmd = [
+                sys.executable,
+                "-m",
+                "mlc_llm",
+                "compile",
+                str(model_path),
+                "--opt",
+                opt,
+                "--overrides",
+                overrides,
+                "--device",
+                device,
+                "--output",
+                dso_path,
+            ]
+            if system_lib_prefix:
+                cmd += ["--system-lib-prefix", system_lib_prefix + "_"]
+            logger.info("Compiling using commands below:")
+            logger.info("%s", blue(shlex.join(cmd)))
+            subprocess.run(cmd, check=False, env=os.environ)
+            # note on windows: compilation can succeed but return code is still nonzero
+            # check whether file exists instead
+            if not os.path.isfile(dso_path):
+                raise RuntimeError("Cannot find compilation output, compilation failed")
+            shutil.move(dso_path, dst)
+            logger.info("Using compiled model lib: %s", bold(dst))
+
+    hash_key = {
+        "model_config": _get_model_config(),
+        "overrides": _get_overrides(),
+        "opt": _get_optimization_flags(),
+        "device": device2str(device) if isinstance(device, Device) else device,
+        "model_type": model_type,
+        "quantization": quantization,
+    }
+    if device in ["iphone", "macabi", "android"]:
+        if system_lib_prefix is None:
+            system_lib_hash_value = hashlib.md5(
+                json.dumps(
+                    hash_key,
+                    sort_keys=True,
+                    indent=2,
+                ).encode("utf-8")
+            ).hexdigest()
+            system_lib_prefix = f"{model_type}_{quantization}_{system_lib_hash_value}".replace(
+                "-", "_"
+            )
+        hash_key["system_lib_prefix"] = system_lib_prefix
+    hash_value = hashlib.md5(
+        json.dumps(
+            hash_key,
+            sort_keys=True,
+            indent=2,
+        ).encode("utf-8")
+    ).hexdigest()
+    dst = MLC_LLM_HOME / "model_lib" / f"{hash_value}.{lib_suffix}"
+    if dst.is_file() and MLC_JIT_POLICY in ["ON", "READONLY"]:
+        logger.info("Using cached model lib: %s", bold(str(dst)))
+        return JITResult(str(dst), system_lib_prefix)
+    if MLC_JIT_POLICY == "READONLY":
+        raise RuntimeError(
+            "No cached model lib found, and JIT is disabled by MLC_JIT_POLICY=READONLY"
+        )
+    _run_jit(
+        opt=hash_key["opt"],
+        overrides=hash_key["overrides"],
+        device=hash_key["device"],
+        system_lib_prefix=system_lib_prefix,
+        dst=str(dst),
+    )
+    return JITResult(str(dst), system_lib_prefix)

python/mlc_llm/interface/package.py

@@ -0,0 +1,402 @@
+"""Python entrypoint of package."""
+
+import dataclasses
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Literal  # noqa: UP035
+
+from mlc_llm.interface import jit
+from mlc_llm.support import download_cache, logging, style
+
+logging.enable_logging()
+logger = logging.getLogger(__name__)
+
+SUPPORTED_DEVICES = ["iphone", "macabi", "android"]
+
+
+def build_model_library(
+    package_config: Dict[str, Any],  # noqa: UP006
+    device: str,
+    bundle_dir: Path,
+    app_config_path: Path,
+) -> Dict[str, str]:  # noqa: UP006
+    """Build model libraries. Return the dictionary of "library prefix to lib path"."""
+    # - Create the bundle directory.
+    os.makedirs(bundle_dir, exist_ok=True)
+    # Clean up all the directories in `output/bundle`.
+    logger.info('Clean up all directories under "%s"', str(bundle_dir))
+    for content_path in bundle_dir.iterdir():
+        if content_path.is_dir():
+            shutil.rmtree(content_path)
+
+    # - Process each model, and prepare the app config.
+    app_config_model_list = []
+
+    model_entries = package_config.get("model_list", [])
+    if not isinstance(model_entries, list):
+        raise ValueError('The "model_list" in "mlc-package-config.json" is expected to be a list.')
+    model_lib_path_for_prepare_libs = package_config.get("model_lib_path_for_prepare_libs", {})
+    if not isinstance(model_lib_path_for_prepare_libs, dict):
+        raise ValueError(
+            'The "model_lib_path_for_prepare_libs" in "mlc-package-config.json" is expected to be '
+            "a dict."
+        )
+
+    jit.log_jit_policy()
+
+    for model_entry in package_config.get("model_list", []):
+        # - Parse model entry.
+        if not isinstance(model_entry, dict):
+            raise ValueError('The element of "model_list" is expected to be a dict.')
+        model = model_entry["model"]
+        model_id = model_entry["model_id"]
+        bundle_weight = model_entry.get("bundle_weight", False)
+        overrides = model_entry.get("overrides", {})
+        model_lib = model_entry.get("model_lib", None)
+
+        estimated_vram_bytes = model_entry["estimated_vram_bytes"]
+        if not isinstance(model, str):
+            raise ValueError('The value of "model" in "model_list" is expected to be a string.')
+        if not isinstance(model_id, str):
+            raise ValueError('The value of "model_id" in "model_list" is expected to be a string.')
+        if not isinstance(bundle_weight, bool):
+            raise ValueError(
+                'The value of "bundle_weight" in "model_list" is expected to be a boolean.'
+            )
+        if not isinstance(overrides, dict):
+            raise ValueError('The value of "overrides" in "model_list" is expected to be a dict.')
+        if model_lib is not None and not isinstance(model_lib, str):
+            raise ValueError('The value of "model_lib" in "model_list" is expected to be string.')
+
+        # - Load model config. Download happens when needed.
+        model_path = download_cache.get_or_download_model(model)
+
+        # - Jit compile if the model lib path is not specified.
+        model_lib_path = (
+            model_lib_path_for_prepare_libs.get(model_lib, None) if model_lib is not None else None
+        )
+        if model_lib_path is None:
+            if model_lib is None:
+                logger.info(
+                    'Model lib is not specified for model "%s". Now jit compile the model library.',
+                    model_id,
+                )
+            else:
+                logger.info(
+                    'Model lib path for "%s" is not specified in "model_lib_path_for_prepare_libs".'
+                    "Now jit compile the model library.",
+                    model_lib,
+                )
+            model_lib_path, model_lib = dataclasses.astuple(
+                jit.jit(
+                    model_path=model_path,
+                    overrides=overrides,
+                    device=device,
+                    system_lib_prefix=model_lib,
+                    skip_log_jit_policy=True,
+                )
+            )
+            assert model_lib is not None
+            model_lib_path_for_prepare_libs[model_lib] = model_lib_path
+
+        # - Set "model_url"/"model_path" and "model_id"
+        app_config_model_entry = {}
+        is_local_model = not model.startswith("HF://") and not model.startswith("https://")
+        app_config_model_entry["model_id"] = model_id
+        app_config_model_entry["model_lib"] = model_lib
+
+        # - Bundle weight
+        if is_local_model and not bundle_weight:
+            raise ValueError(
+                f'Model "{model}" in "model_list" is a local path.'
+                f'Please set \'"bundle_weight": true\' in the entry of model "{model}".'
+            )
+        if bundle_weight:
+            if not os.path.isfile(model_path / "tensor-cache.json"):
+                raise ValueError(
+                    f'Bundle weight is set for model "{model}". However, model weights are not'
+                    f'found under the directory "{model}". '
+                    + (
+                        "Please follow https://llm.mlc.ai/docs/compilation/convert_weights.html to "
+                        "convert model weights."
+                        if is_local_model
+                        else "Please report this issue to https://github.com/mlc-ai/mlc-llm/issues."
+                    )
+                )
+            # Overwrite the model weight directory in bundle.
+            bundle_model_weight_path = bundle_dir / model_id
+            logger.info(
+                "Bundle weight for %s, copy into %s",
+                style.bold(model_id),
+                style.bold(str(bundle_model_weight_path)),
+            )
+            if bundle_model_weight_path.exists():
+                shutil.rmtree(bundle_model_weight_path)
+            shutil.copytree(model_path, bundle_model_weight_path)
+        if bundle_weight and device in ["iphone", "macabi"]:
+            app_config_model_entry["model_path"] = model_id
+        else:
+            app_config_model_entry["model_url"] = model.replace("HF://", "https://huggingface.co/")
+
+        # - estimated_vram_bytes
+        app_config_model_entry["estimated_vram_bytes"] = estimated_vram_bytes
+
+        app_config_model_list.append(app_config_model_entry)
+
+    # - Dump "mlc-app-config.json".
+    app_config_json_str = json.dumps(
+        {"model_list": app_config_model_list},
+        indent=2,
+    )
+    with open(app_config_path, "w", encoding="utf-8") as file:
+        print(app_config_json_str, file=file)
+        logger.info(
+            'Dump the app config below to "%s":\n%s',
+            str(app_config_path),
+            style.green(app_config_json_str),
+        )
+    return model_lib_path_for_prepare_libs
+
+
+def validate_model_lib(
+    app_config_path: Path,
+    package_config_path: Path,
+    model_lib_path_for_prepare_libs: dict,
+    device: Literal["iphone", "macabi", "android"],
+    output: Path,
+) -> None:
+    """Validate the model lib prefixes of model libraries."""
+    if device == "android":
+        from tvm.support import ndk as cc
+    else:
+        from tvm.support import cc
+
+    with open(app_config_path, encoding="utf-8") as file:
+        app_config = json.load(file)
+
+    tar_list = []
+    model_set = set()
+
+    for model, model_lib_path in model_lib_path_for_prepare_libs.items():
+        model_lib_path = os.path.join(model_lib_path)
+        lib_path_valid = os.path.isfile(model_lib_path)
+        if not lib_path_valid:
+            raise RuntimeError(f"Cannot find file {model_lib_path} as an {device} model library")
+        tar_list.append(model_lib_path)
+        model_set.add(model)
+
+    os.makedirs(output / "lib", exist_ok=True)
+    if device in ["iphone", "macabi"]:
+        lib_name = "libmodel_iphone.a"
+    else:
+        lib_name = "libmodel_android.a"
+    lib_path = output / "lib" / lib_name
+
+    def _get_model_libs(lib_path: Path) -> List[str]:  # noqa: UP006
+        """Get the model lib prefixes in the given static lib path."""
+        global_symbol_map = cc.get_global_symbol_section_map(lib_path)
+        libs = []
+        suffix = "___tvm_ffi__library_bin"
+        for name, _ in global_symbol_map.items():
+            if name.endswith(suffix):
+                model_lib = name[: -len(suffix)]
+                if model_lib.startswith("_"):
+                    model_lib = model_lib[1:]
+                libs.append(model_lib)
+        return libs
+
+    cc.create_staticlib(lib_path, tar_list)
+    available_model_libs = _get_model_libs(lib_path)
+    logger.info("Creating lib from %s", str(tar_list))
+    logger.info("Validating the library %s", str(lib_path))
+    logger.info(
+        "List of available model libs packaged: %s,"
+        " if we have '-' in the model_lib string, it will be turned into '_'",
+        str(available_model_libs),
+    )
+    global_symbol_map = cc.get_global_symbol_section_map(lib_path)
+    error_happened = False
+
+    for item in app_config["model_list"]:
+        model_lib = item["model_lib"]
+        model_id = item["model_id"]
+        if model_lib not in model_set:
+            # NOTE: this cannot happen under new setting
+            # since if model_lib is not included, it will be jitted
+            raise RuntimeError(
+                f"ValidationError: model_lib={model_lib} specified for model_id={model_id} "
+                "is not included in model_lib_path_for_prepare_libs argument, "
+                "This will cause the specific model not being able to load, "
+                f"model_lib_path_for_prepare_libs={model_lib_path_for_prepare_libs}"
+            )
+
+        model_prefix_pattern = model_lib.replace("-", "_") + "___tvm_ffi__library_bin"
+        if (
+            model_prefix_pattern not in global_symbol_map
+            and "_" + model_prefix_pattern not in global_symbol_map
+        ):
+            # NOTE: no lazy format is ok since this is a slow pass
+            model_lib_path = model_lib_path_for_prepare_libs[model_lib]
+            log_msg = (
+                "ValidationError:\n"
+                f"\tmodel_lib {model_lib} requested in {str(app_config_path)}"
+                f" is not found in {str(lib_path)}\n"
+                f"\tspecifically the model_lib for {model_lib_path}.\n"
+                f"\tcurrent available model_libs in {str(lib_path)}: {available_model_libs}\n"
+                f"\tThis can happen when we manually specified model_lib_path_for_prepare_libs"
+                f" in {str(package_config_path)}\n"
+                f"\tConsider remove model_lib_path_for_prepare_libs (so library can be jitted)"
+                "or check the compile command"
+            )
+            logger.info(log_msg)
+            error_happened = True
+
+    if not error_happened:
+        logger.info(style.green("Validation pass"))
+    else:
+        logger.info(style.red("Validation failed"))
+        sys.exit(255)
+
+
+def build_android_binding(mlc_llm_source_dir: Path, output: Path) -> None:
+    """Build android binding in MLC LLM"""
+    mlc4j_path = mlc_llm_source_dir / "android" / "mlc4j"
+
+    # Move the model libraries to "build/lib/" for linking
+    os.makedirs(Path("build") / "lib", exist_ok=True)
+    src_path = str(output / "lib" / "libmodel_android.a")
+    dst_path = str(Path("build") / "lib" / "libmodel_android.a")
+    logger.info('Moving "%s" to "%s"', src_path, dst_path)
+    shutil.move(src_path, dst_path)
+
+    # Build mlc4j
+    logger.info("Building mlc4j")
+    subprocess.run([sys.executable, mlc4j_path / "prepare_libs.py"], check=True, env=os.environ)
+    # Copy built files back to output directory.
+    lib_path = output / "lib" / "mlc4j"
+    os.makedirs(lib_path, exist_ok=True)
+    logger.info('Clean up all directories under "%s"', str(lib_path))
+    for content_path in lib_path.iterdir():
+        if content_path.is_dir():
+            shutil.rmtree(content_path)
+
+    src_path = str(mlc4j_path / "src")
+    dst_path = str(lib_path / "src")
+    logger.info('Copying "%s" to "%s"', src_path, dst_path)
+    shutil.copytree(src_path, dst_path)
+
+    src_path = str(mlc4j_path / "build.gradle")
+    dst_path = str(lib_path / "build.gradle")
+    logger.info('Copying "%s" to "%s"', src_path, dst_path)
+    shutil.copy(src_path, dst_path)
+
+    src_path = str(Path("build") / "output")
+    dst_path = str(lib_path / "output")
+    logger.info('Copying "%s" to "%s"', src_path, dst_path)
+    shutil.copytree(src_path, dst_path)
+
+    os.makedirs(lib_path / "src" / "main" / "assets")
+    src_path = str(output / "bundle" / "mlc-app-config.json")
+    dst_path = str(lib_path / "src" / "main" / "assets" / "mlc-app-config.json")
+    logger.info('Moving "%s" to "%s"', src_path, dst_path)
+    shutil.move(src_path, dst_path)
+
+
+def build_iphone_binding(mlc_llm_source_dir: Path, output: Path) -> None:
+    """Build iOS binding in MLC LLM"""
+    # Build iphone binding
+    logger.info("Build iphone binding")
+    subprocess.run(
+        ["bash", mlc_llm_source_dir / "ios" / "prepare_libs.sh"],
+        check=True,
+        env=os.environ,
+    )
+
+    # Copy built libraries back to output directory.
+    for static_library in (Path("build") / "lib").iterdir():
+        dst_path = str(output / "lib" / static_library.name)
+        logger.info('Copying "%s" to "%s"', static_library, dst_path)
+        shutil.copy(static_library, dst_path)
+
+
+def build_macabi_binding(mlc_llm_source_dir: Path, output: Path) -> None:
+    """Build Mac Catalyst binding in MLC LLM"""
+    deployment_target = os.environ.get("MLC_MACABI_DEPLOYMENT_TARGET", "18.0")
+    macabi_arch = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64"
+    logger.info("Build macabi binding (deployment target %s)", deployment_target)
+    cmd = [
+        "bash",
+        str(mlc_llm_source_dir / "ios" / "prepare_libs.sh"),
+        "--catalyst",
+        "--deployment-target",
+        deployment_target,
+    ]
+    if macabi_arch:
+        cmd += ["--arch", macabi_arch]
+    subprocess.run(cmd, check=True, env=os.environ)
+
+    # Copy built libraries back to output directory.
+    build_dir = Path(f"build-maccatalyst-{macabi_arch}")
+    for static_library in (build_dir / "lib").iterdir():
+        dst_path = str(output / "lib" / static_library.name)
+        logger.info('Copying "%s" to "%s"', static_library, dst_path)
+        shutil.copy(static_library, dst_path)
+
+
+def package(
+    package_config_path: Path,
+    mlc_llm_source_dir: Path,
+    output: Path,
+) -> None:
+    """Python entrypoint of package."""
+    logger.info('MLC LLM HOME: "%s"', mlc_llm_source_dir)
+
+    # - Read package config.
+    with open(package_config_path, encoding="utf-8") as file:
+        package_config = json.load(file)
+    if not isinstance(package_config, dict):
+        raise ValueError(
+            "The content of MLC package config is expected to be a dict with "
+            f'field "model_list". However, the content of "{package_config_path}" is not a dict.'
+        )
+
+    # - Read device.
+    if "device" not in package_config:
+        raise ValueError(f'JSON file "{package_config_path}" is required to have field "device".')
+    device = package_config["device"]
+    if device not in SUPPORTED_DEVICES:
+        raise ValueError(
+            f'The "device" field of JSON file {package_config_path} is expected to be one of '
+            f'{SUPPORTED_DEVICES}, while "{device}" is given in the JSON.'
+        )
+
+    bundle_dir = output / "bundle"
+    app_config_path = bundle_dir / "mlc-app-config.json"
+    # - Build model libraries.
+    model_lib_path_for_prepare_libs = build_model_library(
+        package_config, device, bundle_dir, app_config_path
+    )
+    # - Validate model libraries.
+    validate_model_lib(
+        app_config_path,
+        package_config_path,
+        model_lib_path_for_prepare_libs,
+        device,
+        output,
+    )
+
+    # - Copy model libraries
+    if device == "android":
+        build_android_binding(mlc_llm_source_dir, output)
+    elif device == "iphone":
+        build_iphone_binding(mlc_llm_source_dir, output)
+    elif device == "macabi":
+        build_macabi_binding(mlc_llm_source_dir, output)
+    else:
+        assert False, "Cannot reach here"
+
+    logger.info("All finished.")

python/mlc_llm/interface/router.py

@@ -0,0 +1,125 @@
+"""Python entrypoint of router."""
+
+from collections.abc import AsyncGenerator
+from http import HTTPStatus
+from typing import List, Literal, Optional, Type  # noqa: UP035
+
+import fastapi
+import uvicorn
+from fastapi.middleware.cors import CORSMiddleware
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.protocol.openai_api_protocol import CompletionLogProbs, CompletionRequest
+from mlc_llm.router import Router
+from mlc_llm.serve import engine_base, engine_utils
+
+
+def serve(
+    model: str,
+    model_lib: Optional[str],
+    router_host: str,
+    router_port: int,
+    endpoint_hosts: List[str],  # noqa: UP006
+    endpoint_ports: List[int],  # noqa: UP006
+    endpoint_num_gpus: List[int],  # noqa: UP006
+    enable_prefix_cache: bool,
+    router_mode: Literal["disagg", "round-robin"] = "round-robin",
+    pd_balance_factor: float = 0.0,
+    router_type: Type[Router] = Router,  # noqa: UP006
+):
+    """Start the router with the specified configuration."""
+    # 1. Instantiate router
+    router = router_type(
+        model=model,
+        model_lib=model_lib,
+        hosts=endpoint_hosts,
+        ports=endpoint_ports,
+        num_gpus=endpoint_num_gpus,
+        enable_prefix_cache=enable_prefix_cache,
+        router_mode=router_mode,
+        pd_balance_factor=pd_balance_factor,
+    )
+
+    router_app = fastapi.APIRouter()
+
+    @router_app.post("/v1/completions")
+    async def request_completion(request: CompletionRequest, raw_request: fastapi.Request):
+        """OpenAI-compatible completion API.
+        API reference: https://platform.openai.com/docs/api-reference/completions/create
+        """
+        if router is None:
+            return error_protocol.create_error_response(
+                HTTPStatus.BAD_REQUEST, message="Router is not initialized."
+            )
+        request_id = f"cmpl-{engine_utils.random_uuid()}"
+
+        # Streaming response.
+        if request.stream:
+            # We manually get the first response from generator to
+            # capture potential exceptions in this scope, rather then
+            # the StreamingResponse scope.
+            stream_generator = router.handle_completion(request, request_id)
+            first_response = await anext(  # noqa: F821
+                stream_generator
+            )
+
+            async def completion_stream_generator() -> AsyncGenerator[str, None]:
+                if isinstance(first_response, StopAsyncIteration):
+                    yield "data: [DONE]\n\n"
+                    return
+                yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
+                async for response in stream_generator:
+                    yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
+                yield "data: [DONE]\n\n"
+
+            return fastapi.responses.StreamingResponse(
+                completion_stream_generator(), media_type="text/event-stream"
+            )
+
+        # FIXME: Non-streaming response not fully implemented
+        request_final_usage = None
+        output_texts = [""] * request.n
+        finish_reasons: List[Optional[str]] = [None] * request.n  # noqa: UP006
+        logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n  # noqa: UP006
+
+        async for response in router.handle_completion(request, request_id):
+            if await raw_request.is_disconnected():
+                # In non-streaming cases, the engine will not be notified
+                # when the request is disconnected.
+                # Therefore, we check if it is disconnected each time,
+                # and explicitly return.
+                # Note that requesta abort is triggered when the async for and funciton scope ends.
+                return error_protocol.create_error_response(
+                    HTTPStatus.BAD_REQUEST, message="The request has disconnected"
+                )
+            # TODO(Charlie): This is copied from engine.py --
+            # why is it here? Non-streaming only has a single chunk right?
+            # this is the final chunk
+            # if response.usage is not None:
+            #     request_final_usage = response.usage
+            #     continue
+            for choice in response.choices:
+                output_texts[choice.index] += choice.text
+                if choice.finish_reason is not None and finish_reasons[choice.index] is None:
+                    finish_reasons[choice.index] = choice.finish_reason
+                if choice.logprobs is not None:
+                    logprob_results[choice.index] = choice.logprobs
+
+        assert all(finish_reason is not None for finish_reason in finish_reasons)
+        return engine_base.wrap_completion_response(
+            request_id=request_id,
+            model=request.model,
+            output_texts=output_texts,
+            finish_reasons=finish_reasons,
+            logprob_results=logprob_results,
+            usage=request_final_usage,
+        )
+
+    # 2. Set up app
+    app = fastapi.FastAPI()
+    app.add_middleware(CORSMiddleware)
+    app.include_router(router_app)
+    app.exception_handler(error_protocol.BadRequestError)(error_protocol.bad_request_error_handler)
+
+    # 3. Run
+    uvicorn.run(app, host=router_host, port=router_port, log_level="info")

python/mlc_llm/interface/serve.py

@@ -0,0 +1,131 @@
+"""Python entrypoint of serve."""
+
+from typing import Any, List, Literal, Optional, Tuple, Union  # noqa: UP035
+
+import fastapi
+import uvicorn
+from fastapi.middleware.cors import CORSMiddleware
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.serve import engine
+from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine
+from mlc_llm.serve.entrypoints import (
+    debug_entrypoints,
+    metrics_entrypoints,
+    microserving_entrypoints,
+    openai_entrypoints,
+)
+from mlc_llm.serve.server import ServerContext
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+def serve(
+    model: str,
+    device: str,
+    model_lib: Optional[str],
+    mode: Literal["local", "interactive", "server"],
+    enable_debug: bool,
+    additional_models: List[Union[str, Tuple[str, str]]],  # noqa: UP006
+    embedding_model: Optional[str],
+    embedding_model_lib: Optional[str],
+    tensor_parallel_shards: Optional[int],
+    pipeline_parallel_stages: Optional[int],
+    opt: Optional[str],
+    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],
+    spec_tree_width: Optional[int],
+    prefix_cache_mode: Literal["disable", "radix"],
+    prefix_cache_max_num_recycling_seqs: Optional[int],
+    prefill_mode: Literal["hybrid", "chunked"],
+    enable_tracing: bool,
+    host: str,
+    port: int,
+    allow_credentials: bool,
+    allow_origins: Any,
+    allow_methods: Any,
+    allow_headers: Any,
+    api_key: Optional[str] = None,
+):
+    """Serve the model with the specified configuration."""
+    # Create engine and start the background loop
+    async_engine = engine.AsyncMLCEngine(
+        model=model,
+        device=device,
+        model_lib=model_lib,
+        mode=mode,
+        engine_config=engine.EngineConfig(
+            additional_models=additional_models,
+            tensor_parallel_shards=tensor_parallel_shards,
+            pipeline_parallel_stages=pipeline_parallel_stages,
+            opt=opt,
+            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,
+            spec_tree_width=spec_tree_width,
+            prefix_cache_mode=prefix_cache_mode,
+            prefix_cache_max_num_recycling_seqs=prefix_cache_max_num_recycling_seqs,
+            prefill_mode=prefill_mode,
+        ),
+        enable_tracing=enable_tracing,
+    )
+
+    # Set up embedding model if specified
+    emb_engine = None
+    if embedding_model is not None:
+        if embedding_model_lib is None:
+            raise ValueError(
+                "--embedding-model-lib is required when --embedding-model is specified."
+            )
+        emb_engine = AsyncEmbeddingEngine(
+            model=embedding_model,
+            model_lib=embedding_model_lib,
+            device=device,
+        )
+        logger.info("Embedding model %s loaded successfully.", embedding_model)
+
+    with ServerContext() as server_context:
+        server_context.add_model(model, async_engine)
+        if emb_engine is not None:
+            server_context.add_embedding_engine(embedding_model, emb_engine)
+        server_context.api_key = api_key
+
+        app = fastapi.FastAPI()
+        app.add_middleware(
+            CORSMiddleware,
+            allow_credentials=allow_credentials,
+            allow_origins=allow_origins,
+            allow_methods=allow_methods,
+            allow_headers=allow_headers,
+        )
+
+        app.include_router(openai_entrypoints.app)
+        app.include_router(metrics_entrypoints.app)
+        app.include_router(microserving_entrypoints.app)
+
+        server_context.enable_debug = enable_debug
+
+        if enable_debug:
+            app.include_router(debug_entrypoints.app)
+            logger.info("Enable debug endpoint and debug_config in requests...")
+
+        app.exception_handler(error_protocol.BadRequestError)(
+            error_protocol.bad_request_error_handler
+        )
+        uvicorn.run(app, host=host, port=port, log_level="info")